scan with gradient checkpointing
Open
- Dominant language
- Python
- Stars
- 36.3k
- Forks
- 3.8k
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 408
Description
It would be great to have a version of lax.scan used a recursive gradient checkpointing (e.g., "binomial checkpointing") that allows for differentiating through long time series with logarithmic time/space costs.
In principle this could be built on top of the experimental `remat` decorator: https://github.com/google/jax/pull/1749
Contributor guide
Assessment
This issue has not been assessed yet.