cache cond/switch branch wrappers to avoid tracing cache misses

flax/nnx/transforms/general.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import functools
 import typing as tp
+import weakref
 
 from flax.nnx import (
   extract,
@@ -158,6 +159,9 @@ def split_inputs_wrapper(*args):
 
   return split_inputs_wrapper  # type: ignore
 
+_merge_inputs_cache: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
+
+
 @tp.overload
 def merge_inputs(
   *,
@@ -192,6 +196,15 @@ def merge_inputs(
   if isinstance(f, Missing):
     return functools.partial(merge_inputs, ctxtag=ctxtag)  # type: ignore[return-value]
 
+  try:
+    inner = _merge_inputs_cache.get(f)
+  except TypeError:
+    inner = None
+  if inner is not None:
+    wrapper = inner.get(ctxtag)
+    if wrapper is not None:
+      return wrapper  # type: ignore
+
   @functools.wraps(f)
   def merge_inputs_wrapper(*pure_args):
     args = extract.from_tree(pure_args, ctxtag=ctxtag, is_inner=True)
@@ -200,4 +213,11 @@ def merge_inputs_wrapper(*pure_args):
     pure_args_out, pure_out = extract.to_tree((args_out, out), ctxtag=ctxtag)
     return pure_args_out, pure_out
 
+  if inner is None:
+    try:
+      inner = {}
+      _merge_inputs_cache[f] = inner
+    except TypeError:
+      return merge_inputs_wrapper  # type: ignore
+  inner[ctxtag] = merge_inputs_wrapper
   return merge_inputs_wrapper  # type: ignore

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
 
@@ -585,6 +586,27 @@ def __call__(self, *operands):
     return out, updates
 
 
+_simple_cond_fn_cache: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
+
+
+def _get_simple_cond_fn(f, graph):
+  try:
+    inner = _simple_cond_fn_cache.get(f)
+  except TypeError:
+    return SimpleCondFn(f, graph=graph)
+  if inner is None:
+    inner = {}
+    try:
+      _simple_cond_fn_cache[f] = inner
+    except TypeError:
+      return SimpleCondFn(f, graph=graph)
+  wrapper = inner.get(graph)
+  if wrapper is None:
+    wrapper = SimpleCondFn(f, graph=graph)
+    inner[graph] = wrapper
+  return wrapper
+
+
 def cond(
   pred,
   true_fun: tp.Callable[..., A],
@@ -622,8 +644,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=graph),
+      _get_simple_cond_fn(false_fun, graph=graph),
       *operands,
     )
     if graph:
@@ -677,7 +699,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=graph) for f in branches],
       *operands,
     )
     if graph:

tests/nnx/transforms_test.py

@@ -7121,6 +7121,42 @@ def false_fn(m):
     with self.assertRaises(ValueError):
       nnx.cond(True, true_fn, false_fn, m, graph=False)
 
+  @parameterized.parameters(True, False)
+  def test_cond_wrapper_identity_stable(self, graph):
+    from flax.nnx.transforms.transforms import _get_simple_cond_fn
+
+    def true_fn(x):
+      return x + 1.0
+
+    def false_fn(x):
+      return x - 1.0
+
+    self.assertIs(
+        _get_simple_cond_fn(true_fn, graph),
+        _get_simple_cond_fn(true_fn, graph),
+    )
+    self.assertIsNot(
+        _get_simple_cond_fn(true_fn, graph),
+        _get_simple_cond_fn(false_fn, graph),
+    )
+    self.assertIsNot(
+        _get_simple_cond_fn(true_fn, True),
+        _get_simple_cond_fn(true_fn, False),
+    )
+
+  def test_cond_repeated_calls_correct(self):
+    def true_fn(x):
+      return x + 1.0
+
+    def false_fn(x):
+      return x - 1.0
+
+    for _ in range(3):
+      np.testing.assert_allclose(
+          nnx.cond(True, true_fn, false_fn, 1.0, graph_updates=False), 2.0)
+      np.testing.assert_allclose(
+          nnx.cond(False, true_fn, false_fn, 1.0, graph_updates=False), 0.0)
+
 class TestSwitch(parameterized.TestCase):
   @parameterized.parameters(
     (True, False),