google-deepmind/tf2jax v0.3.6

tf2jax 0.3.6

Repository: google-deepmind/tf2jax

Tag: v0.3.6

Published: 2024-03-26T10:37:59Z

Prerelease: no

Release notes:

What's Changed

• Add Cumprod to tf2jax ops. https://github.com/google-deepmind/tf2jax/pull/149
• Fix squeeze in tf2jax for empty tuple axis arguments. https://github.com/google-deepmind/tf2jax/pull/150
• Update expected attrs of VarHandleOp -- add debug_name. https://github.com/google-deepmind/tf2jax/pull/154
• Only enable shape assertions during refinement if all input shapes are static. https://github.com/google-deepmind/tf2jax/pull/155
• Refactor is_poly_dim to avoid depending on hard-coded module path. https://github.com/google-deepmind/tf2jax/pull/156
• Avoid reuse of polymorphic variable names in tf2jax internal. https://github.com/google-deepmind/tf2jax/pull/158
• Update deepmind -> google-deepmind. https://github.com/google-deepmind/tf2jax/pull/159
• Change conversion of library functions and gradient functions to be lazy. https://github.com/google-deepmind/tf2jax/pull/162
• Fix gradient function lookup error when higher order gradient functions are referenced in a jax2tf model but not serialized. https://github.com/google-deepmind/tf2jax/pull/164
• Enforce autograph=False on gradient function. https://github.com/google-deepmind/tf2jax/pull/172
• Update chex and haiku version in https://github.com/google-deepmind/tf2jax/pull/175
• Fix shape inference for jax2tf vjp functions. https://github.com/google-deepmind/tf2jax/pull/180
• Support native serialization with multiple target platforms. https://github.com/google-deepmind/tf2jax/pull/183

Full Changelog: https://github.com/google-deepmind/tf2jax/compare/v0.3.5...v0.3.6