ndarray.at.set(mode="drop") gives incorrect value for the last element in the array about jax HOT 4 OPEN

Hi - I think this is working as expected. In three of the cases, you are setting indices_are_sorted and/or unique_indices to True, when your indices are neither unique nor sorted. If you lie to the compiler, you get undefined results.

In the last case where both flags are set to False, the output is also consistent with the documented behavior of the API: you have multiple duplicate indices. From the lax.scatter documentation:

If multiple updates are performed to the same index of operand, they may be applied in any order.

[2, 3, 4, 7] is a valid output here, because you specified that index -1 may be set to 7.

from jax.

Ah, I was expecting that negatives indices would be always out-of-bounds and that out-of-bounds indices would not be counted against indices_are_sorted and unique_indices. It seems the way to do this, then, would be to use jnp.arange(a.size) + a.size as out of bounds indices, then set indices_are_sorted=False unique_indices=True?

from jax.

It seems the way to do this, then, would be to use jnp.arange(a.size) + a.size as out of bounds indices

Yes, I believe that would work correctly.

then set indices_are_sorted=False unique_indices=True?

I'm not sure whether out-of-bound indices are included or excluded from the compiler's specialized logic around unique or sorted indices. The safest thing would be to not set either of these flags to True, unless your indices are actually unique and sorted (including any out-of-bound indices).

from jax.

@ezhulenev do out-of-bounds indices need to be sorted and unique for XLA Scatter ops in FILL_OR_DROP mode to set indices_are_sorted() && unique_indices()?

On XLA:GPU indices_are_sorted appears to be unused as far as I can tell: GitHub link. So the performance impact of having to set indices_are_sorted=False is none.

... while unique_indices allows scatters to use non-atomic store instructions: GitHub link. So, it's important that we guarantee this with the jnp.arange trick.

from jax.