Stackless yashful

PiperOrigin-RevId: 681582933
google · Oct 3, 2024 · 641bda9 · 641bda9
1 parent 67e2917
commit 641bda9
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()