Remove chex dependencies.

Author: Nush395

colabs/fine_tuning_tutorial_jax.ipynb

@@ -60,11 +60,11 @@
         "# @title Python imports\n",
         "import pathlib\n",
         "from typing import Any, Mapping, Iterator\n",
+        "import dataclasses\n",
         "import enum\n",
         "import functools\n",
         "\n",
         "# We import JAX and some related packages.\n",
-        "import chex\n",
         "import jax\n",
         "import jax.numpy as jnp\n",
         "import optax\n",
@@ -363,7 +363,8 @@
       },
       "outputs": [],
       "source": [
-        "@chex.dataclass(frozen=True)\n",
+        "@jax.tree_util.register_dataclass\n",
+        "@dataclasses.dataclass(frozen=True)\n",
         "class TrainingInput:\n",
         "  # Input tokens given to the model\n",
         "  input_tokens: jax.Array\n",

recurrentgemma/jax/sampler_test.py

@@ -18,9 +18,9 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import chex
 import jax
 import jax.numpy as jnp
+import numpy as np
 from recurrentgemma import common
 import recurrentgemma.jax as griffin_lib
 
@@ -229,12 +229,12 @@ def test_forward_equivalence(self, dtype: str):
         return_logits=True,
     )
     total_sampled_tokens = total_generation_steps + token_input.shape[-1]
-    chex.assert_shape(
-        jnp.array(output_sampler.logits),
+    self.assertEqual(
+        jnp.asarray(output_sampler.logits).shape,
         (batch_size, total_sampled_tokens, model_config.vocab_size),
     )
-    chex.assert_shape(
-        jnp.array(output_sampler.tokens),
+    self.assertEqual(
+        jnp.asarray(output_sampler.tokens).shape,
         (batch_size, total_sampled_tokens),
     )
 
@@ -248,7 +248,7 @@ def test_forward_equivalence(self, dtype: str):
       rtol = 1e-6
       atol = 1e-6
 
-    chex.assert_trees_all_close(
+    np.testing.assert_allclose(
         output_forward,
         out_logits,
         rtol=rtol,