Merge nnx.errors to flax.errors

flax/errors.py

@@ -64,6 +64,15 @@ def __reduce__(self):
    return (FlaxError, (str(self),))
 
 
+#################################################
+# NNX errors                                    #
+#################################################
+
+
+class TraceContextError(FlaxError):
+  pass
+
+
 #################################################
 # lazy_init.py errors                           #
 #################################################

flax/nnx/object.py

@@ -25,13 +25,13 @@
 import numpy as np
 
 from flax.nnx import (
-  errors,
   reprlib,
   tracers,
 )
 from flax.nnx import graph
 from flax.nnx.variables import Variable, VariableState
 from flax.typing import Key
+from flax import errors
 
 G = tp.TypeVar('G', bound='Object')
 

flax/nnx/variables.py

@@ -22,7 +22,7 @@
 
 import jax
 
-from flax import nnx
+from flax import errors
 from flax.nnx import reprlib, tracers
 from flax.typing import Missing
 import jax.tree_util as jtu
@@ -259,7 +259,7 @@ def __setattr__(self, name: str, value: Any) -> None:
 
   def _setattr(self, name: str, value: tp.Any):
     if not self._trace_state.is_valid():
-      raise nnx.errors.TraceContextError(
+      raise errors.TraceContextError(
         f'Cannot mutate {type(self).__name__} from a different trace level'
       )
 

tests/nnx/module_test.py

@@ -17,7 +17,7 @@
 from typing import Any, TypeVar
 
 from absl.testing import absltest
-from flax import nnx
+from flax import nnx, errors
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -39,7 +39,7 @@ def test_trace_level(self):
     @jax.jit
     def f():
       with self.assertRaisesRegex(
-          nnx.errors.TraceContextError,
+          errors.TraceContextError,
           "Cannot mutate 'Dict' from different trace level",
       ):
         m.a = 2

tests/nnx/rngs_test.py

@@ -21,6 +21,7 @@
 from absl.testing import absltest
 
 from flax import nnx
+from flax import errors
 
 
 class TestRngs(absltest.TestCase):
@@ -58,7 +59,7 @@ def test_rng_trace_level_constraints(self):
     @jax.jit
     def f():
       with self.assertRaisesRegex(
-          nnx.errors.TraceContextError,
+          errors.TraceContextError,
           'Cannot call RngStream from a different trace level',
       ):
         rngs.params()
@@ -76,7 +77,7 @@ def h():
 
     self.assertIsInstance(rngs1, nnx.Rngs)
     with self.assertRaisesRegex(
-      nnx.errors.TraceContextError,
+      errors.TraceContextError,
       'Cannot call RngStream from a different trace level',
     ):
       rngs1.params()