jax
b99958db - Places the `remat` decorator on top of the body function.

Commit

2 years ago

Places the `remat` decorator on top of the body function. PiperOrigin-RevId: 577320028

Author

Yunlong Liu

Committer

a-googler

Parents