.at.set(, indices_are_sorted=True/False) for 2D indices

[PhilipVinc]

What does it mean for 2D indices to be sorted?
The documentation says simply that https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html

indices_are_sorted (bool) – If True, the implementation will assume that the indices passed to at[] are sorted in ascending order, which can lead to more efficient execution on some backends.

But it's unclear what it means for higher dimensional indices.
Should I understand it as if the memory offset in the underlying array (essentially, the index in the flattened version of a Nd array) are sorted?

[dfm]

I'm not 100% sure (maybe someone else who knows more will weight in), but you might take a look at the XLA docs for more details about the expected behavior because this parameter is just passed through to XLA: https://openxla.org/xla/operation_semantics

[jakevdp]

Agreed – JAX does nothing with indices_are_sorted beyond passing it to the underlying XLA gather or scatter operation. I'd look at the generated jaxpr and compare it to the requirements in the XLA documentation to figure out the meaning in this context (I'd answer less mysteriously, but the meaning of the indices_are_sorted parameter is not something I've ever fully understood either).