What's Changed
- Update broken Oryx github links to point to spun-off Oryx github repo. by @copybara-service in #67
- Add
Layout
support tojax.jit
. by @copybara-service in #72 - Add 'cond_clobber' mode that lets different loop iterations sow different values to be reaped. by @copybara-service in #66
- Fixes monkey patching attribute error triggered by cl/629129623. by @copybara-service in #74
- Replace deprecated
jax.tree_*
functions withjax.tree.*
by @copybara-service in #76
Full Changelog: v0.2.6...v0.2.7