From 7b63edca200acedc4bfbf725c92e23d245ee9eb3 Mon Sep 17 00:00:00 2001 From: Dougal Maclaurin Date: Wed, 2 Oct 2024 14:01:18 -0700 Subject: [PATCH] Stackless yashful PiperOrigin-RevId: 681582933 --- flax/core/tracers.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/flax/core/tracers.py b/flax/core/tracers.py index fe2ff874c0..9d4393e967 100644 --- 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()