Update xla_sharding import path to new location

We are moving the TensorFlow APIs outside of XLA and will remove the old
path soon.

PiperOrigin-RevId: 477859491

Author: lingvo-bot

lingvo/core/gshard_layers.py

@@ -28,7 +28,12 @@
 
 # pylint: disable=g-direct-tensorflow-import
 from tensorflow.compiler.tf2xla.python import xla
-from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding
+# pylint: disable=g-import-not-at-top
+try:
+  from tensorflow.python.compiler.xla.experimental import xla_sharding
+except ImportError:
+  # OSS backward compatibility, can be removed when TF is updated.
+  from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding
 # pylint: enable=g-direct-tensorflow-import
 
 Split = gshard_utils.Split

lingvo/core/gshard_utils.py

@@ -25,7 +25,12 @@
 import sentencepiece as sentencepiece_processor
 # pylint: disable=g-direct-tensorflow-import
 from tensorflow.compiler.xla import xla_data_pb2
-from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding
+# pylint: disable=g-import-not-at-top
+try:
+  from tensorflow.python.compiler.xla.experimental import xla_sharding
+except ImportError:
+  # OSS backward compatibility, can be removed when TF is updated.
+  from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding
 # pylint: enable=g-direct-tensorflow-import
 
 ThreadLocalStack = thread_local_utils.ThreadLocalStack

lingvo/core/var_tmp_wrappers.py

@@ -22,7 +22,12 @@
 from lingvo import compat as tf
 
 # pylint: disable=g-direct-tensorflow-import
-from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding
+# pylint: disable=g-import-not-at-top
+try:
+  from tensorflow.python.compiler.xla.experimental import xla_sharding
+except ImportError:
+  # OSS backward compatibility, can be removed when TF is updated.
+  from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding
 # pylint: enable=g-direct-tensorflow-import