Persistent Compilation Cache - explicit storing inner function

[JeyRunner]

Is there a way to tell the jax compilation cache (https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html) to store a specific jitten function?

I have something like:

@jit
def inner_fn(value):
  """ takes very long to compile """
  # ...

@jit(static_arg_names=['params'])
def outer_fn(value, params)
   result = inner_fn(value + ...)

outer(some_values, some_params)

The compilation cache seems to just save outer_fn on disk, but I also want inner_fn to be saved as well since inner_fn will always stay the same, but the static arguments 'params' may change from run to run. So I don't want inner_fn to be recompiled every time outer_fn changes (its params or implementation).
Note that inner_fn is very slow to compile (uses mjx).

[patrick-kidger]

I don't believe so. Everything inside of the top-level JIT is inlined and so the inner JIT decorator doesn't actually do anything.

FWIW there is also eqx.internal.noinline which allows you to avoid recompiling inner_fn every time just outer_fn changes. Heads up that this does come up the multiple dubious banners of "wildly experimental", "this might only work on the CPU", and "I have no idea how this interacts with the compilation cache"... !

[JeyRunner]

Thank you, I'll give eqx.internal.noinline a try and check if it works in combination with the compilation cache :)
Do you think this may have a negative effect on performance?

[patrick-kidger]

It does have a substantial negative impact on runtime performance, I'm afraid.
Ultimately this is something that JAX just really isn't optimised around!

[JeyRunner]

Ok, this would make the benefit of faster compilation time kind of useless :/
Hopefully, there will be some solution for this in the future (without runtime performance disadvantages).

Maybe I should open an issue, but no idea how this could be implemented.
(maybe just dumping the tracing results to disk and then recompile without retracing, but the compile time advantages with this would probably be rather small...)