TRACING CACHE MISS at flax/core/scope.py:950:18 (Scope.param)

[DBraun]

I have a large autogenerated Flax Linen Module that is not working well with caching. I'm using

• flax 0.10.5
• jaxlib 0.5.3
• jax 0.5.3

on WSL2 w/ NVIDIA RTX 4080 Super driver 572.70 (newer drivers cause exit code 139 but that's a different story...)

cache_issue.py:

import jax

jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
jax.config.update("jax_pprint_use_color", True)
jax.config.update("jax_explain_cache_misses", True)

from jax import numpy as jnp, random
from algo15 import DX7Algo15, FAUSTFLOAT


if __name__ == "__main__":

    sample_rate = 44_100

    model = DX7Algo15(sample_rate=sample_rate)

    key = random.key(0)

    duration_sec = 1.

    N_SAMPLES = int(duration_sec * sample_rate)
    N_CHANNELS = model.num_inputs

    x = jnp.ones((N_CHANNELS, N_SAMPLES), dtype=FAUSTFLOAT)
    x = x.at[0, :].set(440)

    unroll = 1

    variables = model.init({"params": key, "rng_stream": key}, x, length=N_SAMPLES, unroll=unroll)

    @jax.jit
    def forward(variables, x: jnp.ndarray):
        y = model.apply(variables, x, length=N_SAMPLES, unroll=unroll, rngs={"rng_stream": key})
        return y

    y = forward(variables, x)

    assert y.ndim == 2
    assert y.shape[0] == model.num_outputs
    assert y.shape[1] == x.shape[1]
    assert y.shape[1] == N_SAMPLES

    print("All done!")

In the first run, the ending is

WARNING:2025-06-22 19:19:07,153:jax._src.compiler:107: PERSISTENT COMPILATION CACHE MISS for 'jit_forward' with key 'jit_forward-277bb252941ee475b89264b7444b78689454bb1ae5ea3536b2179087a81f7aaa'
WARNING:2025-06-22 19:19:17,873:jax._src.compilation_cache:265: Writing jit_forward to persistent compilation cache with key 'jit_forward-277bb252941ee475b89264b7444b78689454bb1ae5ea3536b2179087a81f7aaa'
All done!

In later runs, the output has hundreds of variations (different lambda id) of

WARNING:2025-06-22 19:23:28,057:jax._src.pjit:1151: TRACING CACHE MISS at /home/braun/.local/lib/python3.12/site-packages/flax/core/scope.py:950:18 (Scope.param) because:
  never seen function:
    Scope.param.<locals>.<lambda> id=138999571240800 defined at /home/braun/.local/lib/python3.12/site-packages/flax/core/scope.py:951
  but seen another function defined on the same line; maybe the function is
  being re-defined repeatedly, preventing caching?

and then

WARNING:2025-06-22 19:23:28,578:jax._src.pjit:1151: TRACING CACHE MISS at /mnt/c/Users/braun/GitHub/myproject/cache_issue.py:38:8 (<module>) because:
  never seen function:
    forward id=139001048154176 defined at /mnt/c/Users/braun/GitHub/myproject/cache_issue.py:33
All done!

Note that the __call__ of DX7Algo15 has a lambda tick:

	def __call__(self, x: jnp.ndarray, length: int = None, unroll: int = 1) -> jnp.ndarray:

		if length is None and x is not None:
			length = x.shape[-1]

		# Handle generators (no input case)
		if x is None:
			x = jnp.zeros((self.num_inputs, length), dtype=FAUSTFLOAT)

		carry = self.initialize_carry()

		# Unnormalize parameters once before the scan
		params = self.unnormalize()

		def tick(module, carry, *xs):
			return module.tick(params, carry, *xs)

		scan_fn = nn.scan(tick,
			variable_broadcast="params",
			split_rngs={"rng_stream": True},
			length=length,
			unroll=unroll,
			in_axes=1,
			out_axes=1,
		)
		new_carry, outputs = scan_fn(self, carry, x)

		return outputs

Can I get some help with improving caching?

cache_issue.py.txt
algo15.py.txt