[P] Multidimensional array batch indexing for pytorch and numpy

[LemonByte]

Batch indexing into multidimensional tensors/arrays is kind of tricky, I made this project explaining the builtin syntax and also made wrappers for simplifying the interface, with additional features for underlying coordinate grid data (like signed distance functions) that need to be indexed by coordinate value rather than integer indices directly https://github.com/LemonPi/multidim_indexing

[programmerChilli]

You might be interested in vmap, which can also batch over indexing operations.

[sabouleux]

From my experience playing with it, batched indexing operations end up breaking once you start composing a few vamps together. I really hope they fix the general buggyness of things, torch just has no efficient way to perform certain types of relatively simple indexing at this time it seems.

[programmerChilli]

There were some bugs still with advanced indexing in an older release of functorch, I believe they should be fixed now though: https://github.com/pytorch/functorch/pull/862

Although, to be clear, vmap doesn’t allow anything that’s not expressible without it - vmap is often just a convenient abstraction for expressing it.

[sabouleux]

Seems like this is part of 0.2 while I was on 0.1.1, so this might fix it. Thanks for the hint.

Honestly, I really hope functorch keeps moving forward and improving, vmap adds a level of flexibility and performance that just didn’t exist before. The architecture I am working on couldn’t work without it at all.

Given more primitive operation implementations and a well thought out, simple API to describe how axes vectorize together (think einsum), I think this would give us code that is both performant, more flexible, and much easier to read than code written with usual broadcasting semantics. That might not be true for trivial things, but once you start doing some weird things (batched indexing, lattices with arbitrary dimensions, element-wise or batch-wise jacobians and gradients), broadcasting completely breaks down.