Merge nnx.errors to flax.errors

google · Sep 10, 2024 · 1b03f4d · 1b03f4d
1 parent f948154
commit 1b03f4d
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 18 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')