eindex for pytorch

[shekhovt]

Hi,

This looks great! I am struggling with indexing every time. Previously I was able to utilize gather to do what was needed, but now I need something like:
out[i,j,k] = in[j,k,idx[i,j]]
which I have no idea how to implement (thinking of gather with expanding both in and idx). Your notation would be so transparent and helpful! Perhaps I missed something, do you already have some functionality for pytorch? Or what do the perspectives look like?

[arogozhnikov]

Thanks for kind words.

Torch some time ago committed to support array api standard, but I don't see much progress on their side recently:
pytorch/pytorch#58743

Whether this should be considered 'won't do' - I don't know, but currently eindex does not work for torch.

[arogozhnikov]

actually, maybe https://github.com/data-apis/array-api-compat can solve the issue, haven't tried

[shekhovt]

Thanks for explaining the situation with standard array api.

but currently eindex does work for torch

Sorry, did you mean then "does NOT work for torch"?

[arogozhnikov]

yes :) mis-edited that sentence

[p-i-]

I was delighted to stumble upon this work while digging through feature requests on einops.

I would love to see eindex supported in PyTorch.

Is it known whether they are interested to implement?

This would facilitate research into gating strategies.

Yesterday I was trying to do what should be a very simple operation: Given a vector of offsets o (0 or 1), compose a y that for each element k chooses y1[k] if o[k] else y0[k]. And to additionally handle a batch dimension.

I was not happy with the look of the result:

        splitter_indices = (λ >= 0).long()  # (nBatch, nHidden)

        # (2, nHidden, nOut)
        splitter_indices_expanded = splitter_indices.unsqueeze(-1).expand(-1, -1, nOut)

        # self.Y: (2, nHidden, nOut)
        selected_Y = torch.gather(self.Y, 0, splitter_indices_expanded)

        scaled_chosen_Ys = λλ.unsqueeze(-1) * selected_Y

It is the opposite of "Pythonic" -- quite impossible for me myself to read without having to apply some unpleasant mental contortion every time I look at it.

I simply want to get self.Y[splitter_indices] over the batch.

So I'd be hoping to write something along the vague lines of: self.Y_k[splitterIndices_k] where the k subscript ranges over the batch dimension.

Or "batch i, M[2 *i] Y -> batch Y", which I'd read as:

• First arg: A batch dimension (preserved in the output) of indices (represented i; we know they're indices because they're used in M[2 *i], using the C lingo of *i meaning "value at location i"
• Second arg: An object where the first two dimensions (2 and i) are represented by M, third is Y
• ->
• preserve the batch dim, each item is Y

I find myself at a point of considering looking at Julia for prototyping designs. I tried playing with it when it came out, and found it to be rather fierce. But this seems to be an intrinsically fierce game. It gets to the point where it's questionable whether to even seek a clean Python implementation, or put up with ugly (and heavily commented) prototypes and write clean C++/CUDA once something works.

I do hope team PyTorch give this library the attention it deserves.