Add numpy.put_along_axis.

[carlosgmartin]

Fixes #9954.

[jakevdp]

Thanks for looking into this! It's a good start, but we'll need to handle (and test!) all the broadcasting semantics of put_along_axis – i.e. both indices and values may be N-dimensional arrays that broadcast together along all axes except axis

[carlosgmartin]

@jakevdp Isn't that already handled? For example, testPutAlongAxis tests the case where M and J are different.

By the way, this seems to be equivalent to

def put_along_axis(arr, indices, values, axis):
  idx = jnp.indices(arr.shape, sparse=True)
  idx = idx[:axis] + (indices,) + idx[axis:][1:]  # replace the indices for axis
  return arr.at[idx].set(values)

[carlosgmartin]

@jakevdp I think it should be fixed now.

[carlosgmartin]

@jakevdp Updated.

[carlosgmartin]

@jakevdp Updated.

[carlosgmartin]

@jakevdp I reduced the number of test cases further. Let me know if you want to keep the utility function jax._src.numpy.util.broadcastable_shapes I created. (Perhaps it might be useful in the future.)

[jakevdp]

Looks good, thanks!

[jakevdp]

It looks like the example output has a stray character

[jakevdp]

We're seeing test failures on H100. I think it's probably due to the fact that repeated indices within scatter may be updated in any order, and so the updated values may not match the updated values in NumPy, which always traverses the indices in order.

Probably the best way to fix this would be to ensure within tests that generated indices are unique within each slice, though I don't know of an easy way to do that in the general case using existing test utilities.