Tupled indexing


I’ve a tensor of the following dimensions (#Batches, 2, 1).

Now, I’ve to perform an operation such that I’ve to apply a function that only operates on some elements of the batches (say, if #batches == 10, let them be divided into 3,3,4).

So, I need a function that has the type signature of

(#batches, 2, 1) -> IndexList -> (#batches, 1)

However, inside the function, it should use the index list and apply a function only on the indices specified (the index list is exhaustive and the different chunks are contiguous in the batch dimension).

How do I go about doing that?