Help understanding start_index_map in jax.lax.gather

[davisyoshida]

Despite staring at the JAX and XLA docs for gather for quite some time, I've been unable to understand some of the arguments. For example, in testing some toy examples to try to build my intuition, I hit a case where errors told me:

• I needed 3 values in start_index_map
• Those values must be in [0, 2)
• The values must all be distinct

There aren't enough holes to hold all these pigeons! Can anyone give me a few minimal examples which illustrate the effect of various values of start_index_map?

[jakevdp]

Is there a reason you need to use lax.gather? My number one recommendation would be to never use lax.gather directly, in which case there's no need to understand what start_index_map means. Everything that can be expressed via a call to gather can be expressed more succinctly and more understandably with standard numpy-style indexing.

But if you for some reason absolutely need to use gather, I don't think anyone can explain it any better than it's explained in the documentation at https://www.tensorflow.org/xla/operation_semantics#general_semantics. Yes, it really is that complicated, thus my initial recommendation 😀

[hawkinsp]

+1 to Jake's comment. I am unable to keep details of gather in my memory. I need to relearn it every time I have to mess with it. If you must learn it, I recommend looking at examples of the jaxpr or HLO that JAX produces from things like x[...] indexing, which is much simpler to think about.

[davisyoshida]

@jakevdp I'm writing a transformation which needs to know how to handle lax.gather, so if I don't understand the semantics of it I'm worried that my transformation will be incorrect.

[davisyoshida]

That's probably too vague, so basically what I'm doing is writing a transformation to apply LoRA to a function without needing to manually modify it. The problem is that matrices aren't just used in matmuls, they're also used in embedding lookups and convs, so my interpreter needs to at the very least handle lax.dot_general, lax.conv_general_dilated, and lax.gather. I may just give up on generality and raise an error if anything which looks different than a simple embedding lookup is done with gather though.