Remove chex dependency

Author: samanklesaria

pyproject.toml

@@ -11,12 +11,7 @@ authors = [
     {name = "Flax team", email = "flax-dev@google.com"},
 ]
 dependencies = [
-    # temporary numpy version fix due to
-    # https://github.com/google/flax/issues/5162
-    # and
-    # https://github.com/google-deepmind/chex/issues/424
-    "numpy>=1.23.2,<2.4.0; python_version>='3.11'",
-    "numpy>=1.26.0,<2.4.0; python_version>='3.12'",
+    "numpy>=1.23.2",
     # keep in sync with jax-version in .github/workflows/build.yml
     "jax>=0.8.1",
     "msgpack",
@@ -43,7 +38,6 @@ dynamic = ["version", "readme"]
 [project.optional-dependencies]
 testing = [
     "clu",
-    "chex",
     "clu<=0.0.9; python_version<'3.10'",
     "einops",
     "gymnasium[atari]; python_version<'3.14'",

tests/jax_utils_test.py

@@ -20,7 +20,6 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import chex
 from flax import jax_utils
 import jax
 import jax.numpy as jnp
@@ -29,13 +28,30 @@
 NDEV = 4
 
 
-class PadShardUnpadTest(chex.TestCase):
+def assert_max_traces(n):
+  """Decorator to assert that a function is traced at most n times."""
+  def decorator(fn):
+    trace_count = {'count': 0}
+
+    def wrapped(*args, **kwargs):
+      trace_count['count'] += 1
+      if trace_count['count'] > n:
+        raise AssertionError(
+          f"Function was traced {trace_count['count']} times, "
+          f"expected at most {n} traces"
+        )
+      return fn(*args, **kwargs)
+
+    wrapped.trace_count = trace_count
+    return wrapped
+  return decorator
+
+
+class PadShardUnpadTest(parameterized.TestCase):
   BATCH_SIZES = [NDEV, NDEV + 1, NDEV - 1, 5 * NDEV, 5 * NDEV + 1, 5 * NDEV - 1]
   DTYPES = [np.float32, np.uint8, jax.numpy.bfloat16, np.int32]
 
-  def tearDown(self):
-    chex.clear_trace_counter()
-    super().tearDown()
+
 
   @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)
   def test_basics(self, dtype, bs):
@@ -47,7 +63,7 @@ def add(a, b):
 
     x = np.arange(bs, dtype=dtype)
     y = add(x, 10 * x)
-    chex.assert_type(y.dtype, x.dtype)
+    self.assertEqual(y.dtype, x.dtype)
     np.testing.assert_allclose(np.float64(y), np.float64(x + 10 * x))
 
   @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)
@@ -59,24 +75,22 @@ def add(a, b):
 
     x = jnp.arange(bs, dtype=dtype)
     y = add(dict(a=x), (10 * x,))
-    chex.assert_type(y.dtype, x.dtype)
+    self.assertEqual(y.dtype, x.dtype)
     np.testing.assert_allclose(np.float64(y), np.float64(x + 10 * x))
 
   @parameterized.parameters(DTYPES)
   def test_min_device_batch_avoids_recompile(self, dtype):
     @partial(jax_utils.pad_shard_unpad, static_argnums=())
     @jax.jit
-    @chex.assert_max_traces(n=1)
+    @assert_max_traces(n=1)
     def add(a, b):
       b = jnp.asarray(b, dtype=dtype)
       return a + b
 
-    chex.clear_trace_counter()
-
     for bs in self.BATCH_SIZES:
       x = jnp.arange(bs, dtype=dtype)
       y = add(x, 10 * x, min_device_batch=9)  # pylint: disable=unexpected-keyword-arg
-      chex.assert_type(y.dtype, x.dtype)
+      self.assertEqual(y.dtype, x.dtype)
       np.testing.assert_allclose(np.float64(y), np.float64(x + 10 * x))
 
   @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)
@@ -87,7 +101,7 @@ def add(a, b):
 
     x = jnp.arange(bs, dtype=dtype)
     y = add(x, 10)
-    chex.assert_type(y.dtype, x.dtype)
+    self.assertEqual(y.dtype, x.dtype)
     np.testing.assert_allclose(np.float64(y), np.float64(x + 10))
 
   @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)
@@ -102,7 +116,7 @@ def add(params, a, *, b):
 
     x = jnp.arange(bs, dtype=dtype)
     y = add(5, x, b=10)
-    chex.assert_type(y.dtype, x.dtype)
+    self.assertEqual(y.dtype, x.dtype)
     np.testing.assert_allclose(np.float64(y), np.float64(5 * x + 10))
 
 

tests/linen/linen_test.py

@@ -19,7 +19,6 @@
 from typing import Any
 
 from absl.testing import absltest, parameterized
-import chex
 from flax import ids
 from flax import linen as nn
 from flax.linen import fp8_ops
@@ -960,7 +959,7 @@ def _strip_partitioning(x):
     np.testing.assert_array_equal(got_scale, expected_scale)
 
     # Compares the rest of PyTree nodes
-    chex.assert_trees_all_close(expected, got)
+    jax.tree.map(lambda x, y: np.testing.assert_allclose(x, y), expected, got)
 
 
 class StochasticTest(parameterized.TestCase):