flax

d539ce78 - Add nnx.make_jaxpr.

Add nnx.make_jaxpr. This change introduces `nnx.make_jaxpr`, a version of `jax.make_jaxpr`.`graph_updates=True` is not supported and will raise an error. PiperOrigin-RevId: 889901203