Merge pull request #4186 from IvyZX:errors

PiperOrigin-RevId: 673510178
google · Sep 11, 2024 · b967964 · b967964
2 parents 61ea8a6 + e8b0a1f
commit b967964
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 24 deletions.
diff --git a/flax/errors.py b/flax/errors.py
@@ -64,6 +64,15 @@ def __reduce__(self):
    return (FlaxError, (str(self),))
 
 
+#################################################
+# NNX errors                                    #
+#################################################
+
+
+class TraceContextError(FlaxError):
+  pass
+
+
 #################################################
 # lazy_init.py errors                           #
 #################################################

diff --git a/flax/nnx/errors.py b/flax/nnx/errors.py
diff --git a/flax/nnx/object.py b/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')
 

diff --git a/flax/nnx/variables.py b/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
@@ -235,7 +235,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'
       )
 

diff --git a/tests/nnx/module_test.py b/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

diff --git a/tests/nnx/rngs_test.py b/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()