CPU backend gather fails with multi-dimensional indices despite documented [N, L] support

[michalharakal]

Summary

TensorOps.gather documents embedding-style lookups with indices of shape [L] or [N, L] (DefaultCpuOps.kt comment: Input: [vocabSize, embeddingDim], Indices: [L] or [N, L], Output: [L, embeddingDim] or [N, L, embeddingDim]), but the CPU implementation reads the indices with a flat single-index access:

val indexList = IntArray(numIndices) { i ->
    val v = indices.data[i]   // ItemsAccessor.get(vararg) validates index count
    (v as Number).toInt()
}

For rank-2 indices this throws IllegalArgumentException: Number of indices (1) must match tensor dimensions (2).

Reproduce (SKaiNET 0.36.0, DirectCpuExecutionContext)

val ctx = DirectCpuExecutionContext()
val table = ctx.fromFloatArray<FP32, Float>(Shape(10, 4), FP32::class, FloatArray(40) { it.toFloat() })
val ids = ctx.fromIntArray<Int32, Int>(Shape(2, 3), Int32::class, intArrayOf(0, 1, 2, 7, 8, 9))
ctx.ops.gather(table, ids as Tensor<DType, *>, dim = 0)
// -> IllegalArgumentException: Number of indices (1) must match tensor dimensions (2)

1-D indices work fine. Workaround: flatten indices to [N*L], gather, reshape the result back to [N, L, embDim].

Suggested fix

Convert the flat loop to multi-dimensional iteration (or read through a row-major flat view of the indices tensor). The rest of the implementation already computes the correct output shape for rank-2 indices. Happy to send a PR.

Found while implementing a batched embedding lookup ([B, T] -> [B, T, D]) on SKaiNET 0.36.0.

[michalharakal]

Fixed by the merged PR (#877 for #863/#864, #879 for #859); the auto-close keyword didn't fire because the PR addressed several issues at once.