mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix jax2tf import so it works with both the latest tensorflow release (2.10.0) and tf-nightly

Browse Source
This commit is contained in:
Skye Wanderman-Milne 2022-09-30 15:55:22 -07:00
parent fb8558cfdd
commit 0a69c9a27b
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
jax/experimental/jax2tf/jax2tf.py
View File

@ -68,7 +68,11 @@ import tensorflow as tf # type: ignore[import]
from tensorflow.compiler.tf2xla.python import xla as tfxla # type: ignore[import]
from tensorflow.compiler.xla import xla_data_pb2 # type: ignore[import]
from tensorflow.core.framework import attr_value_pb2 # type: ignore[import]
from tensorflow.python.compiler.xla.experimental import xla_sharding # type: ignore[import]
try:
from tensorflow.python.compiler.xla.experimental import xla_sharding # type: ignore[import]
except ModuleNotFoundError:
# This can be removed when TF 2.10 support is no longer needed.
from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding # type: ignore[import]
from tensorflow.python.framework import ops as tf_ops # type: ignore[import]
# pylint: enable=g-direct-tensorflow-import