Stackless yashful

PiperOrigin-RevId: 681582933
google · Oct 4, 2024 · 9d1094e · 9d1094e
1 parent 6469c5d
commit 9d1094e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/flax/core/tracers.py b/flax/core/tracers.py
@@ -31,6 +31,7 @@ def current_trace():
   return jax.core.get_opaque_trace_state(convention="flax")
 
 def check_trace_level(base_level):
-  level = current_trace()
-  if level != base_level:
-    raise errors.JaxTransformError()
+  pass
+  # level = current_trace()
+  # if level != base_level:
+  #   raise errors.JaxTransformError()