Tf2jax

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

What's Changed
* Better error message for missing inputs in https://github.com/deepmind/tf2jax/pull/131
* Support sharding annotation in https://github.com/deepmind/tf2jax/pull/132, https://github.com/deepmind/tf2jax/pull/133
* Support shape refinement with XlaCallModule in https://github.com/deepmind/tf2jax/pull/140, in https://github.com/deepmind/tf2jax/pull/144, in https://github.com/deepmind/tf2jax/pull/145
* Add LowerBound and UpperBound operation to tf2jax to support searchsorted. in https://github.com/deepmind/tf2jax/pull/141
* Dropped support for python 3.8 in https://github.com/deepmind/tf2jax/pull/140,
* Handle shape check asserts in in https://github.com/deepmind/tf2jax/pull/146

**Full Changelog**: https://github.com/deepmind/tf2jax/compare/v0.3.4...v0.3.5

What's Changed
* Add documentation regarding jax2tf native serialization and XlaCallModule. https://github.com/deepmind/tf2jax/pull/123
* Improving support for XlaCallModule. https://github.com/deepmind/tf2jax/pull/97, https://github.com/deepmind/tf2jax/pull/122, https://github.com/deepmind/tf2jax/pull/118, https://github.com/deepmind/tf2jax/pull/100
* Support FFT, FFT2D, FFT3D, IFFT, IFFT2D, IFFT3D, RFFT2D, RFFT3D, IRFFT2D, IRFFT3D. https://github.com/deepmind/tf2jax/pull/115
* Support Bessel0e, Bessel1e and PopulationCount. https://github.com/deepmind/tf2jax/pull/114
* Support LeakyReluOp. https://github.com/deepmind/tf2jax/pull/120
* Support BatchToSpaceND, SpaceToBatchND. by https://github.com/deepmind/tf2jax/pull/126
* Support MatrixSetDiagV3. https://github.com/deepmind/tf2jax/pull/111
* Handle VarHandleOp . https://github.com/deepmind/tf2jax/pull/105
* Add placeholder for XlaSharding. https://github.com/deepmind/tf2jax/pull/103
* Better error message for missing parameters. https://github.com/deepmind/tf2jax/pull/110
* Better handling of static arguments for ScatterND. https://github.com/deepmind/tf2jax/pull/112
* Improved handling of non-tensor outputs. https://github.com/deepmind/tf2jax/pull/98


**Full Changelog**: https://github.com/deepmind/tf2jax/compare/v0.3.3...v0.3.4

What's Changed
* Make TF2JAX run on M1 macs. https://github.com/deepmind/tf2jax/pull/81
* Use input_signature when available. https://github.com/deepmind/tf2jax/pull/83
* Preemptive fix for upcoming JAX changes. https://github.com/deepmind/tf2jax/pull/85, https://github.com/deepmind/tf2jax/pull/86
* Preemptive fix for polymorphic shape support. https://github.com/deepmind/tf2jax/pull/88
* Re-enable native lowering test for nested custom gradient. https://github.com/deepmind/tf2jax/pull/91
* Support XlaReducePrecision. https://github.com/deepmind/tf2jax/pull/93


**Full Changelog**: https://github.com/deepmind/tf2jax/compare/v0.3.2...v0.3.3

What's Changed
* Support XlaOptimizationBarrier. https://github.com/deepmind/tf2jax/pull/63
* Support TensorListFromTensor, TensorListGetItem, TensorListSetItem, TensorListReserve, TensorListStack. https://github.com/deepmind/tf2jax/pull/78, https://github.com/deepmind/tf2jax/pull/79
* Support DvNoNan, ResourceGather, ResizeNearestNeighbor, RFFT, IRFFT, Angle, StatelessRandomGetAlg, EnsureShape, CheckNumerics, Roll. https://github.com/deepmind/tf2jax/pull/60, https://github.com/deepmind/tf2jax/pull/61, https://github.com/deepmind/tf2jax/pull/62, https://github.com/deepmind/tf2jax/pull/64, https://github.com/deepmind/tf2jax/pull/67, https://github.com/deepmind/tf2jax/pull/68, https://github.com/deepmind/tf2jax/pull/74
* Support experimental XlaCallModule used by jax2tf. https://github.com/deepmind/tf2jax/pull/70, https://github.com/deepmind/tf2jax/pull/71
* Add heuristic for inferring cumulative reductions and avoid potentially expensive reduce_window calls. https://github.com/deepmind/tf2jax/pull/55
* Fix an edge case in StrideSlice. https://github.com/deepmind/tf2jax/pull/80
* Dropped support for python 3.7 (to match JAX). https://github.com/deepmind/tf2jax/pull/73

**Full Changelog**: https://github.com/deepmind/tf2jax/compare/v0.3.1...v0.3.2

What's Changed
* Improved documentation. https://github.com/deepmind/tf2jax/pull/49, https://github.com/deepmind/tf2jax/pull/50
* Improved handling of input specs. https://github.com/deepmind/tf2jax/pull/48
* Support Cholesky, Eig, SelfAdjointEigV2, SVD, QR, MatrixTriangularSolve. https://github.com/deepmind/tf2jax/pull/45, https://github.com/deepmind/tf2jax/pull/46, https://github.com/deepmind/tf2jax/pull/47, https://github.com/deepmind/tf2jax/pull/51
* Support UnsortedSegmentMax, UnsortedSegmentMin, UnsortedSegmentProd, UnsortedSegmentSum and TensorScatterUpdate. https://github.com/deepmind/tf2jax/pull/52


**Full Changelog**: https://github.com/deepmind/tf2jax/compare/v0.3.0...v0.3.1