Repository: google/jax

Issue:
OverflowError on assignment of valid value to uint32 array
### Description

```python
import jax.numpy as jnp

a = jnp.zeros((1), dtype=jnp.uint32)
val = 2**32 - 1
# val and its uint32 representation are identical
print(val, jnp.uint32(val))
# assignment works with an explicit cast to uint32
b = a.at[1].set(jnp.uint32(val))
# assignment with an implicit cast throws an exception
b = a.at[1].set(val) # returns OverflowError 
```

### What jax/jaxlib version are you using?

_No response_

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
Thanks for reporting this, and the clear repro!

#15506 might fix it, but it currently causes a "deprecation warning" test to fail (i.e. we no longer raise a warning), and so I need advice from @jakevdp on how to proceed.
(Unlikely to be relevant, but I just noticed that I wrote `a.at[1]` while `a.at[0]` would have made more sense, since `a` is length 1.)
@mattjj - is this something we can fix?

At the JIT boundary, Python scalars like `1` have to be converted to objects of specific dtypes, and we made the decision long ago to not use value-dependent semantics when choosing that dtype. So, unless `jax_enable_x64` is set to True, every Python integer will be represented as an `int32`, regardless of whether it eventually is intended to be converted to `uint32`.

With that in mind, isn't this working as expected?
I should spell-out the subtext of my comment. We want
```python
x.at[0].set(1)
```
to act the same inside and outside JIT, so it follows that this should produce the same result:
```python
jax.jit(x.at[0].set)(1)
```
But when `1` hits a `jit` boundary, it must be converted to a JAX array with a particular `dtype`. We always choose `int32`, and so passing a python scalar that does not fit into `int32` should result in an error.
Thanks for spelling it out clearly! If I understand correctly, I think consistency across the jit boundary is a separate issue from whether assignment works without a jit boundary. Especially, I think assignment of uint32 is already handled differently when it crosses a jit boundary. See this snippet:

```Python
def gen_a(val):
  a = jnp.zeros((1), dtype=jnp.uint32)
  return a.at[0].set(jnp.uint32(val))

val = 2**32 - 1
gen_a(val) # executes
jit(gen_a)(val) # throws OverflowError
```



Yeah, that's a good point. We should be able to support this kind of assignment in the case that the scalar does not cross a tracing boundary.

Resolving patch:
diff --git a/jax/_src/ops/scatter.py b/jax/_src/ops/scatter.py
--- a/jax/_src/ops/scatter.py
+++ b/jax/_src/ops/scatter.py
@@ -64,9 +64,13 @@ def _scatter_update(x, idx, y, scatter_op, indices_are_sorted,
   Returns:
     An ndarray representing an updated `x` after performing the scatter-update.
   """
-
   x = jnp.asarray(x)
-  y = jnp.asarray(y)
+  if (isinstance(y, int) and np.issubdtype(x.dtype, np.integer) and
+      np.iinfo(x.dtype).min <= y <= np.iinfo(x.dtype).max):
+    y = jnp.asarray(y, dtype=x.dtype)
+  else:
+    y = jnp.asarray(y)
+
   # XLA gathers and scatters are very similar in structure; the scatter logic
   # is more or less a transpose of the gather equivalent.
   treedef, static_idx, dynamic_idx = jnp._split_index_for_jit(idx, x.shape)


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -1467,5 +1467,16 @@ def testIndexedUpdateAliasingBug(self):
     y = jnp.zeros(8)
     self.assertArraysEqual(fn(y), jax.jit(fn)(y))
 
+  def testScatterValuesCastToTargetDType(self):
+    # https://github.com/google/jax/issues/15505
+    a = jnp.zeros(1, dtype=jnp.uint32)
+    val = 2**32 - 1  # too large for int32
+
+    b = a.at[0].set(jnp.uint32(val))
+    self.assertEqual(int(b[0]), val)
+
+    c = a.at[0].set(val)
+    self.assertEqual(int(c[0]), val)
+
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())