lowering / cost analysis of @nnx.jit functions

[lzanini]

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]

Converted this into an issue: #4094