flax

9a1bc4eb - Add tree-mode-only nnx.{vjp,jvp,jit_partial}