jax

a5c2c472 - [jax2tf] Added support for x64 for the remaining test files (#4282)

[jax2tf] Added support for x64 for the remaining test files (#4282) * [jax2tf] Added support for x64 in other test files. This includes: - control_flow_ops_test.py - jax2tf_test.py - saved_model_test.py - stax_test