jax
61703690
- Add a direct HLO lowering of remat_p that doesn't call eval_jaxpr.
Go
Login via GitHub
Home
Pricing
FAQ
Install
Login
via GitHub
Commit
View On
GitHub
Commit
1 year ago
Add a direct HLO lowering of remat_p that doesn't call eval_jaxpr. This turns out to be faster, not least because we don't need to use the tracing machinery again. PiperOrigin-RevId: 647462045
References
#22156 - Add a direct HLO lowering of remat_p that doesn't call eval_jaxpr.
Author
hawkinsp
Committer
a-googler
Parents
af8bdd15
Loading