Implement `at` method for `_IndexUpdateRef`

Currently the following does not work

from jax import numpy as jnp

a = jnp.arange(3)
a = a.at[1:].at[:-1].set(0)

though IMHO the desired instruction is clear. The equivalent numpy expression a[1:][:-1] = 0 works according to one’s intuition and sets the second element to 0.

[shoyer]

The ndindex library can do these sort of calculations with as_subindex: https://quansight-labs.github.io/ndindex/api.html#ndindex.ndindex.NDIndex.as_subindex

[jakevdp]

And even for numpy, this kind of composability doesn’t always work as you might expect. For example:

In [1]: import numpy as np

In [2]: idx = np.array([2, 4, 5])

In [3]: x = np.arange(10)

In [4]: x[idx][:3] = 999

In [5]: print(x)
[0 1 2 3 4 5 6 7 8 9]

The reason is that while x[idx] = 999 calls x.__setitem__, x[idx] without an associated assignment calls __getitem__, which in this case returns a copy rather than a view, and the second index expression modifies that copy instead of the original array. So if we were to add such composition, we’d have to think carefully about when and where to mimic / diverge from numpy.