Intermediate value capture API via JAX's `hijax.Box`

[IvyZX]

JAX has this new hijax.Box mechanism that can insert arbitrary values during forward & backward passes. This can be a good alternative to the current sow and perturb APIs on Flax level.

The feature is not complete yet (e.g., not yet work with vmap/scan, and in some other corner cases). Gonna do some prototying first.