Intermediate value capture API via JAX's `hijax.Box`
Open
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
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.
Contributor guide
Assessment
This issue has not been assessed yet.