jax.ops.index_update of int array silently promotes nan to large integer

jax.ops.index_update of int array silently promotes nan to a large integer (i.e. machine capacity for 32 bit int). Is this desired behavior ? The analogous op in numpy throws a ValueError instead.

import jax
import jax.numpy as np
import numpy as onp

a = np.array([0, 47, 0])
b = a.astype(float)

print(jax.ops.index_update(a, a == 0, np.nan)) # [-2147483648          47 -2147483648]
print(jax.ops.index_update(b, b == 0, np.nan)) # [nan 47. nan]

oa = onp.array(a)
oa[oa == 0] = np.nan # ValueError: cannot convert float NaN to integer

[jakevdp]

As of #10924, this behavior emits a warning, and in a future release it will be an error:

>>> import jax.numpy as jnp                                                                                                                                                                                                                                                                                 
>>> jnp.zeros(10, dtype=int).at[3].set(4.5)                                                                                                                                                                                                                                                                 
jax/jax/_src/ops/scatter.py:87: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=float32 to dtype=int32. In future JAX releases this will result in an error.
  warnings.warn("scatter inputs have incompatible types: cannot safely cast "
DeviceArray([0, 0, 0, 4, 0, 0, 0, 0, 0, 0], dtype=int32)

Same with the original issue:

>>> jnp.array([0, 47, 0]).at[0].set(jnp.nan)                                                                                                                                                                                                                                                                
/jax/jax/_src/ops/scatter.py:87: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=float32 to dtype=int32. In future JAX releases this will result in an error.
  warnings.warn("scatter inputs have incompatible types: cannot safely cast "
DeviceArray([ 0, 47,  0], dtype=int32)

[hawkinsp]

We were recently discussing a similar issue:

In [5]: x = jnp.zeros((10), np.int32)

In [6]: x.at[3].set(4.5)
Out[6]: DeviceArray([0, 0, 0, 4, 0, 0, 0, 0, 0, 0], dtype=int32)

This is surprising, and we think it would be better to report an error rather than silently converting from float to int, which is the same problem as in this issue.