flax.nnx.cond causes tracing cache misses

[futscdav]

Using nnx.cond causes tracing cache misses (and thus prevents persistent compilation cache loads) due to creation of transient SimpleCondFn passed to jax.lax.cond. This is very inconvenient when using nnx.cond with larger modules that take a while to compile.

TRACING CACHE MISS at .../flax/nnx/transforms/transforms.py:623:19 (cond):
  never seen function:
    true_fn id=17864732047504 defined at .../repro_cond_miss.py:15
  but seen another function defined on the same line; maybe the function is
  being re-defined repeatedly, preventing caching?
W0622 10:13:43.770388 1182679 partial_eval.py:2164] TRACING CACHE MISS at .../flax/nnx/transforms/transforms.py:623:19 (cond):
  never seen function:
    false_fn id=17864737580928 defined at .../repro_cond_miss.py:19
  but seen another function defined on the same line; maybe the function is
  being re-defined repeatedly, preventing caching?
W0622 10:13:43.771588 1182679 dispatch.py:204] Finished tracing cond for jit in 0.000418901 sec

Here's a repro code that prints the cache misses (including the jax-level prints)

import weakref
from absl import app
from flax import nnx
from flax.nnx.transforms import transforms as nnx_transforms
import jax

jax.config.update("jax_log_compiles", True)
jax.config.update("jax_explain_cache_misses", True)


def true_fn(x):
  return x + 1.0


def false_fn(x):
  return x - 1.0


original_SimpleCondFn = nnx_transforms.SimpleCondFn


def run_test(name, patch_fn):
  print(f"\n--- Running Test: {name} ---", flush=True)

  patch_fn()

  def test_fn(pred, x):
    # Pass graph_updates=False to force the path that uses SimpleCondFn
    return nnx.cond(pred, true_fn, false_fn, x, graph_updates=False)

  for i in range(3):
    print(f"\nIteration {i}:", flush=True)
    pred = jax.numpy.array(i % 2 == 0)
    res = test_fn(pred, 1.0)
    jax.block_until_ready(res)


# Buggy implementation (original behavior, but with prints)
def apply_buggy_patch():
  # We want to intercept calls to SimpleCondFn to print when they are created
  def buggy_SimpleCondFn_factory(f, graph):
    instance = original_SimpleCondFn(f, graph)
    print(
        f"  [BUGGY] Created new wrapper for {f.__name__} (id={id(f)}):"
        f" wrapper_id={id(instance)}",
        flush=True,
    )
    return instance

  nnx_transforms.SimpleCondFn = buggy_SimpleCondFn_factory


# Patched implementation (using cache)
_repro_cache = weakref.WeakKeyDictionary()


def apply_patched_patch():
  _repro_cache.clear()  # Clear cache for clean run

  def patched_SimpleCondFn_factory(f, graph):
    if f not in _repro_cache:
      _repro_cache[f] = {}
    if graph not in _repro_cache[f]:
      instance = original_SimpleCondFn(f, graph)
      _repro_cache[f][graph] = instance
      print(
          f"  [PATCHED] Cache MISS for {f.__name__} (id={id(f)}). Created"
          f" wrapper_id={id(instance)}",
          flush=True,
      )
    else:
      instance = _repro_cache[f][graph]
      print(
          f"  [PATCHED] Cache HIT for {f.__name__} (id={id(f)}). Reused"
          f" wrapper_id={id(instance)}",
          flush=True,
      )
    return instance

  nnx_transforms.SimpleCondFn = patched_SimpleCondFn_factory


def main(argv):
  if len(argv) > 1:
    raise app.UsageError("Too many command-line arguments.")

  print(f"Global true_fn id: {id(true_fn)}")
  print(f"Global false_fn id: {id(false_fn)}")

  # Run buggy first
  run_test("BUGGY (Original)", apply_buggy_patch)

  # Run patched
  run_test("PATCHED", apply_patched_patch)
  print("\n--- Test Finished ---", flush=True)


if __name__ == "__main__":
  app.run(main)

[chenkuanliao]

I've been looking into this and reproduced the cache misses on current main (flax 0.12.7 / jax 0.10.1) with the provided script. The cause is as described: on the graph_updates=False path, nnx.cond builds a fresh SimpleCondFn wrapper for each branch on every call, and since jax.lax.cond keys its tracing cache on callable identity (and SimpleCondFn is a @dataclass(eq=False), so identity-based), each call looks like a brand-new function and re-traces — which also blocks the persistent compilation cache. I confirmed that memoizing the wrappers, as in the repro, turns the per-iteration TRACING CACHE MISS into cache hits. I'm planning to fix this by adding a module-level WeakKeyDictionary keyed by (branch_fn, graph) plus a small helper, and using it in both cond (transforms.py:623) and switch (transforms.py:680) so they share the fix. I'll use weak keys so branch functions aren't kept alive, with a fallback to constructing a fresh wrapper for functions that aren't weak-referenceable (e.g. functools.partial), and add regression tests asserting wrapper-identity stability across repeated calls. I'll open a PR shortly.

[samanklesaria]

Currently reviewing @chenkuanliao's PR, as they were the first to respond to this issue.

[chenkuanliao]

@samanklesaria thanks for the review and comments. I have updated the PR according to your suggestions.