Fix nnx.cond tracing cache misses by caching SimpleCondFn wrappers

flax/nnx/transforms/iteration.py

@@ -17,6 +17,7 @@
 import dataclasses
 import functools
 import typing as tp
+import weakref
 
 
 from flax import struct
@@ -30,6 +31,7 @@
 from flax.nnx.transforms.transforms import (
   resolve_kwargs,
   _resolve_bound_callable,
+  _get_cached_wrapper,
   _raise_bound_method_error,
 )
 from flax.typing import Leaf, Missing, PytreeDeque
@@ -1911,6 +1913,14 @@ def __call__(self, val):
     return self.f(val)
 
 
+_SIMPLE_WHILE_LOOP_BODY_FN_CACHE: weakref.WeakKeyDictionary = (
+  weakref.WeakKeyDictionary()
+)
+_SIMPLE_WHILE_LOOP_COND_FN_CACHE: weakref.WeakKeyDictionary = (
+  weakref.WeakKeyDictionary()
+)
+
+
 @dataclasses.dataclass(eq=False)
 class WhileLoopCondFn:
   f: tp.Callable[..., tp.Any]
@@ -2050,8 +2060,18 @@ def while_loop(cond_fun: tp.Callable[[T], tp.Any],
   if graph_updates is None:
     graph_updates = graphlib.set_graph_updates.current_value()
   if not graph or not graph_updates:
-    simple_body_fn = SimpleWhileLoopBodyFn(body_fun, graph=graph)
-    simple_cond_fn = SimpleWhileLoopCondFn(cond_fun, graph=graph)
+    simple_body_fn = _get_cached_wrapper(
+      _SIMPLE_WHILE_LOOP_BODY_FN_CACHE,
+      SimpleWhileLoopBodyFn,
+      body_fun,
+      graph,
+    )
+    simple_cond_fn = _get_cached_wrapper(
+      _SIMPLE_WHILE_LOOP_COND_FN_CACHE,
+      SimpleWhileLoopCondFn,
+      cond_fun,
+      graph,
+    )
 
     if graph:
       init_val = extract.to_tree2(init_val)
@@ -2094,6 +2114,11 @@ def __call__(self, i, val):
     return out
 
 
+_SIMPLE_FORI_LOOP_BODY_FN_CACHE: weakref.WeakKeyDictionary = (
+  weakref.WeakKeyDictionary()
+)
+
+
 @dataclasses.dataclass(eq=False)
 class ForiLoopBodyFn:
   f: tp.Callable[..., tp.Any]
@@ -2166,7 +2191,12 @@ def fori_loop(lower: int, upper: int,
   if graph_updates is None:
     graph_updates = graphlib.set_graph_updates.current_value()
   if not graph or not graph_updates:
-    simple_body_fn = SimpleForiLoopBodyFn(body_fun, graph=graph)
+    simple_body_fn = _get_cached_wrapper(
+      _SIMPLE_FORI_LOOP_BODY_FN_CACHE,
+      SimpleForiLoopBodyFn,
+      body_fun,
+      graph,
+    )
 
     if graph:
       init_val = extract.to_tree2(init_val)

flax/nnx/transforms/transforms.py

@@ -19,6 +19,7 @@
 import functools
 import inspect
 import typing as tp
+import weakref
 
 from jax._src import checkify as checkify_lib
 
@@ -47,6 +48,7 @@
 M = tp.TypeVar('M', bound=Module)
 MA = tp.TypeVar('MA', bound=Module)
 N = tp.TypeVar('N', bound=Module)
+W = tp.TypeVar('W')
 StrInt = tp.TypeVar('StrInt', str, int)
 AxisName = tp.Hashable
 Leaves = list[Leaf]
@@ -585,6 +587,56 @@ def __call__(self, *operands):
     return out, updates
 
 
+_SIMPLE_COND_FN_CACHE: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
+
+
+def _get_cached_wrapper(
+  cache: weakref.WeakKeyDictionary,
+  wrapper_type: tp.Callable[..., W],
+  f: tp.Callable[..., tp.Any],
+  graph: bool,
+) -> W:
+  """Returns a cached wrapper for ``(f, graph)``.
+
+  Only weak-referenceable, hashable callables can be cached; other callables
+  fall back to a freshly constructed wrapper.
+  """
+  try:
+    weakref.ref(f)
+    hash(f)
+  except TypeError:
+    return wrapper_type(f, graph=graph)
+
+  by_graph = cache.get(f)
+  if by_graph is None:
+    by_graph = {}
+    cache[f] = by_graph
+  wrapper = by_graph.get(graph)
+  if wrapper is None:
+    wrapper = wrapper_type(f, graph=graph)
+    by_graph[graph] = wrapper
+  return wrapper
+
+
+def _get_simple_cond_fn(f: tp.Callable[..., tp.Any], graph: bool) -> SimpleCondFn:
+  """Returns a cached ``SimpleCondFn`` for ``(f, graph)``.
+
+  ``jax.lax.cond`` / ``jax.lax.switch`` key their tracing cache on the identity
+  of the callables they trace. Constructing a fresh ``SimpleCondFn`` on every
+  ``cond``/``switch`` call gives each branch a new identity, defeating that
+  cache and forcing re-tracing (and blocking the persistent compilation cache);
+  see https://github.com/google/flax/issues/5512. Reusing the same wrapper for
+  a given ``(f, graph)`` keeps a stable identity so repeated calls hit the
+  cache.
+
+  ``f`` is held weakly so locally-defined branches stay collectable. Only
+  weak-referenceable, hashable callables can be cached. Callable instances whose
+  classes define ``__slots__`` must include ``'__weakref__'`` in ``__slots__``;
+  otherwise this function falls back to constructing a fresh wrapper.
+  """
+  return _get_cached_wrapper(_SIMPLE_COND_FN_CACHE, SimpleCondFn, f, graph)
+
+
 def cond(
   pred,
   true_fun: tp.Callable[..., A],
@@ -622,8 +674,8 @@ def cond(
     variables = extract.check_no_aliases('cond', operands=operands)
     out, updates = jax.lax.cond(
       pred,
-      SimpleCondFn(true_fun, graph=graph),
-      SimpleCondFn(false_fun, graph=graph),
+      _get_simple_cond_fn(true_fun, graph),
+      _get_simple_cond_fn(false_fun, graph),
       *operands,
     )
     if graph:
@@ -677,7 +729,7 @@ def switch(
     variables = extract.check_no_aliases('switch', operands=operands)
     out, updates = jax.lax.switch(
       index,
-      [SimpleCondFn(f, graph=graph) for f in branches],
+      [_get_simple_cond_fn(f, graph) for f in branches],
       *operands,
     )
     if graph:

tests/nnx/transforms_test.py

@@ -7121,6 +7121,78 @@ def false_fn(m):
     with self.assertRaises(ValueError):
       nnx.cond(True, true_fn, false_fn, m, graph=False)
 
+  def test_get_simple_cond_fn_reuses_wrapper(self):
+    from flax.nnx.transforms import transforms as nnx_transforms
+
+    def f(x):
+      return x
+
+    w1 = nnx_transforms._get_simple_cond_fn(f, True)
+    self.assertIs(w1, nnx_transforms._get_simple_cond_fn(f, True))
+    self.assertIsNot(w1, nnx_transforms._get_simple_cond_fn(f, False))
+    self.assertIsNot(w1, nnx_transforms._get_simple_cond_fn(lambda x: x, True))
+
+  def test_get_simple_cond_fn_non_weakref_fallback(self):
+    from flax.nnx.transforms import transforms as nnx_transforms
+
+    class Slotted:
+      __slots__ = ()
+
+      def __call__(self, x):
+        return x
+
+    f = Slotted()
+    w1 = nnx_transforms._get_simple_cond_fn(f, True)
+    w2 = nnx_transforms._get_simple_cond_fn(f, True)
+    self.assertIsNot(w1, w2)
+
+  @parameterized.named_parameters(
+    ('cond', 'cond'),
+    ('switch', 'switch'),
+    ('while_loop', 'while_loop'),
+    ('fori_loop', 'fori_loop'),
+  )
+  def test_control_flow_reuses_branch_wrappers(self, control_flow):
+    trace_count = 0
+
+    def branch(*args):
+      nonlocal trace_count
+      trace_count += 1
+      return jnp.logical_not(args[-1])
+
+    init_value = jnp.array(False)
+    expected_trace_count = 1
+    if control_flow == 'while_loop':
+      expected_trace_count = 2
+
+    for _ in range(2):
+      if control_flow == 'cond':
+        nnx.cond(
+          jnp.array(True), branch, branch, init_value,
+          graph=False, graph_updates=False,
+        )
+
+      elif control_flow == 'switch':
+        nnx.switch(
+          jnp.array(0), (branch, branch), init_value,
+          graph=False, graph_updates=False,
+        )
+
+      elif control_flow == 'while_loop':
+        nnx.while_loop(
+          branch, branch, init_value, graph=False, graph_updates=False
+        )
+
+      elif control_flow == 'fori_loop':
+        nnx.fori_loop(
+          0, 2, branch, init_value, graph=False, graph_updates=False
+        )
+
+      else:
+        raise AssertionError(f'Unknown control flow: {control_flow}')
+      self.assertEqual(trace_count, expected_trace_count)
+
+
 class TestSwitch(parameterized.TestCase):
   @parameterized.parameters(
     (True, False),