lowering / cost analysis of @nnx.jit functions

[cgarciae]

Discussed in #4093

^{Originally posted by lzanini July 19, 2024}
hi, is there a way to use jax's cost analysis api for nnx.jitted functions ? typically, a training step with a module, an optimizer and a metric. it is unclear to me how to lower the inner jax.jitted function.

[cgarciae]

Currently the inner jitted function is store in an .inner attribute.

f = nnx.jit(lambda x: x)

print(f.inner)  # <PjitFunction of <function jit_fn at ...>>

However, there is no easy way to lower with the exact same inputs that NNX will pass to it. We should provide some way to do this.

[gabbard]

We have an internal implementation of this which works by adding a .lower to the return from jit which follows the same pattern as the existing inner jit_wrapper function to construct the arguments to pass to the .lower of jitted_fn. Would a PR using that approach to this issue be entertained?

[cgarciae]

@gabbard sure! Happy to review a PR.