[Feature Request] Indexing CuArray with CuArray

[Roger-luo]

I guess this can be faster with GPU? But any idea on how to implement this?

A = cu(rand(2048, 2048))
B = CuArray(rand(1:2048, 1024))
A[B, 2]

[MikeInnes]

It should be quite straightforward: the kernel just looks something like C[i] = A[B[i], 2]. Writing a very general kernel is harder though and I guess it'd have to extend GPUArray's current indexing kernels.

[Roger-luo]

Yeah, I agree, I wrote my custom kernel which is not general. I found that pytorch have a general implementation however, so I think we should have one as well.

[maleadt]

This works now.