Replace XLACompatibleSharding with jax.sharding.Sharding since th…

…e former is deprecated and will be removed in the future. PiperOrigin-RevId: 640572753
google · Jun 5, 2024 · 21231af · 21231af
1 parent 617a6cd
commit 21231af
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flax/nnx/nnx/transforms.py b/flax/nnx/nnx/transforms.py
@@ -219,7 +219,7 @@ def jit(
       if the sharding cannot be inferred.
 
       The valid resource assignment specifications are:
-        - :py:class:`XLACompatibleSharding`, which will decide how the value
+        - :py:class:`Sharding`, which will decide how the value
             will be partitioned. With this, using a mesh context manager is not
             required.
         - :py:obj:`None`, will give JAX the freedom to choose whatever sharding