Finalize deprecation of XLACompatibleSharding

PiperOrigin-RevId: 681156145

CHANGELOG.md

@@ -65,6 +65,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     is only a tree-prefix of itself. To preserve the current behavior, you can
     ask `jax.tree.map` to treat `None` as a leaf value by writing:
     `jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None)`.
+  * `jax.sharding.XLACompatibleSharding` has been removed. Please use
+    `jax.sharding.Sharding`.
 
 * Bug fixes
   * Fixed a bug where {func}`jax.numpy.cumsum` would produce incorrect outputs

jax/sharding.py

@@ -31,21 +31,16 @@
 from jax._src.mesh import AbstractMesh
 
 _deprecations = {
-    # Added Jun 4, 2024.
+    # Finalized 2024-10-01; remove after 2025-01-01.
     "XLACompatibleSharding": (
         (
-            "jax.sharding.XLACompatibleSharding is deprecated. Use"
-            " jax.sharding.Sharding instead."
+            "jax.sharding.XLACompatibleSharding was removed in JAX v0.4.34. "
+            "Use jax.sharding.Sharding instead."
         ),
-        _deprecated_XLACompatibleSharding,
+        None,
     )
 }
 
-import typing
-if typing.TYPE_CHECKING:
-  XLACompatibleSharding = _deprecated_XLACompatibleSharding
-else:
-  from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
-  __getattr__ = _deprecation_getattr(__name__, _deprecations)
-  del _deprecation_getattr
-del typing
+from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
+__getattr__ = _deprecation_getattr(__name__, _deprecations)
+del _deprecation_getattr