[JAX] Replace reference to jax._src.lib.xla_client.SingleDeviceSharding with jax.sharding.SingleDeviceSharding, which is its public name.

PiperOrigin-RevId: 826203493

Author: hawkinsp

tests/distillation/distillation_trainer_test.py

@@ -242,7 +242,7 @@ def test_distributed_training(self):
 
     self.assertIsInstance(
         unsharded_variables.layers[0].w1.kernel.value.sharding,
-        jax._src.lib.xla_client.SingleDeviceSharding,
+        jax.sharding.SingleDeviceSharding,
     )
     jax.tree.map_with_path(tc.assert_close, variables, unsharded_variables)
 

tests/sft/checkpoint_manager_test.py

@@ -174,11 +174,11 @@ def test_restore_different_sharding(self):
     # Check the model shardings are restored correctly.
     self.assertIsInstance(
         unsharded_variables.w1.kernel.value.sharding,
-        jax._src.lib.xla_client.SingleDeviceSharding,
+        jax.sharding.SingleDeviceSharding,
     )
     self.assertIsInstance(
         unsharded_variables.w2.kernel.value.sharding,
-        jax._src.lib.xla_client.SingleDeviceSharding,
+        jax.sharding.SingleDeviceSharding,
     )
 
     # Restore the model with shardings.

tests/sft/peft_trainer_test.py

@@ -309,7 +309,7 @@ def test_dist_training(self):
     unsharded_variables = nnx.state(unsharded_model, nnx.Param)
     self.assertIsInstance(
         unsharded_variables.layers[0].w1.kernel.value.sharding,
-        jax._src.lib.xla_client.SingleDeviceSharding,
+        jax.sharding.SingleDeviceSharding,
     )
     jax.tree.map_with_path(tc.assert_close, variables, unsharded_variables)