Merge pull request #1092 from carlosgmartin:safe_increment_unsigned

PiperOrigin-RevId: 683647626

optax/_src/numerics.py

@@ -130,20 +130,15 @@ def safe_increment(count: chex.Numeric) -> chex.Numeric:
   count_dtype = jnp.asarray(count).dtype
   if jnp.issubdtype(count_dtype, jnp.integer):
     max_value = jnp.iinfo(count_dtype).max
-    if jnp.issubdtype(count_dtype, jnp.unsignedinteger):
-      # The comparison count < max_value appears to convert its arguments into
-      # signed integers so we get overflow errors with unsigned integers.
-      raise ValueError(
-          f'Unsigned integers like {count_dtype} cannot be incremented safely.'
-      )
   elif jnp.issubdtype(count_dtype, jnp.floating):
     max_value = jnp.finfo(count_dtype).max
   else:
     raise ValueError(
         f'Cannot safely increment count with dtype {count_dtype},'
         ' valid dtypes are subdtypes of "jnp.integer" or "jnp.floating".'
     )
-  one = jnp.array(1, dtype=count_dtype)
+  max_value = jnp.array(max_value, count_dtype)
+  one = jnp.array(1, count_dtype)
   return jnp.where(count < max_value, count + one, max_value)
 
 

optax/_src/numerics_test.py

@@ -53,10 +53,14 @@ class NumericsTest(chex.TestCase):
       "int16",
       "int32",
       "int64",
+      "uint8",
+      "uint16",
+      "uint32",
+      "uint64",
   ))
   def test_safe_increment(self, dtype):
     """Tests that safe_increment works for all dtypes."""
-    if dtype in ["float64", "int64"]:
+    if dtype in ["float64", "int64", "uint64"]:
       jax.config.update("jax_enable_x64", True)
     dtype = jnp.dtype(dtype)
     inc_fn = self.variant(numerics.safe_increment)
@@ -77,14 +81,11 @@ def test_safe_increment(self, dtype):
       base = jnp.asarray(max_val, dtype=dtype)
       incremented = inc_fn(base)
       np.testing.assert_array_equal(incremented, base)
-    if dtype in ["float64", "int64"]:
+    if dtype in ["float64", "int64", "uint64"]:
       jax.config.update("jax_enable_x64", False)
 
   @parameterized.product(
       str_dtype=[
-          "uint8",
-          "uint16",
-          "uint32",
           "bool",
           "complex64",
       ]