The value ne12 is used in the CUDA code, but I think the indices are being calculated incorrectly. In the CPU code:

const int64_t i12 = i03%ne12;
const int64_t i11 = i02%ne11;
const int64_t i10 = i;

In the CUDA code:

const int i10 = blockIdx.x;
const int i11 = blockIdx.z / ne12; // gridDim.z == ne11*ne12
const int i12 = blockIdx.z % ne12;

In the CUDA code the same values are used for i11/i01 and i12/i02.

Ok, I didn't look in the implementation and assumed it was not implemented. So, will update the PR to fix implementation.

In the CUDA code the same values are used for i11/i01 and i12/i02.

The intention of the operator is that i10 queries rows from src0, hence it corresponds to i01. Respectively:

i10 -> i01
i11 -> i02
i12 -> i03

So I think the CPU implementation is correct. Looking into this.

The CUDA implementation is correct. The problem is that in one of the new GET_ROWS tests, the number of blocks along the 3rd dimension of the kernel exceeds 65536:

Here ne11*n12 > 2^16 and it causes the kernel launch to fail.

For now, I updated the support_op condition to bail out in such cases. Will leave it to you to add proper support for larger sizes.