EnzymeAD / EnzymeAD/Enzyme-JAX
Add Revolve checkpointing
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
Opening an issue to track this. I have this link to the current implementation of square root checkpointing: https://github.com/EnzymeAD/Enzyme-JAX/blob/6c6972ab9808eda4a1905bed3a7ccd013c518041/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp#L558
The goal is to add Revolve checkpointing. @Pangoraw, I'll likely ping you once I have time; however, drop any additional information you have here.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp at the referenced square root checkpointing implementation. Read the surrounding checkpointing logic and the issue discussion to determine the required Revolve behavior and integration details. Done means Revolve checkpointing is implemented and its behavior is validated in the project’s existing checks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100