flax
44b07d02 - fix partial eval for custom_jvp and omnistaging changes

Commit

5 years ago

fix partial eval for custom_jvp and omnistaging changes PiperOrigin-RevId: 325171311

References

Author

levskaya

Committer

a-googler

Parents