jax
2c128ca1 - Make sure out_cts passed to remat ad.backward_pass are to_ct_aval()

Commit

66 days ago

Make sure out_cts passed to remat ad.backward_pass are to_ct_aval() PiperOrigin-RevId: 920010225

References

Author

yashk2810

Committer

Google-ML-Automation

Parents