Hi,
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?