Repository: google/jax
Pull request #21499 (2024-05-29 20:52:28)

Rename `is_compatible_aval` to `check_compatible_aval` since it returns None and not a `bool`.
Rename `is_compatible_aval` to `check_compatible_aval` since it returns None and not a `bool`.



Patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
index 3c4cfcea4..8f7d0a240 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2423,8 +2423,8 @@ def _infer_src_sharding(src, x) -> Sharding | None:
   return None
 
 
-# TODO(yashkatariya): Generalize is_compatible_aval (maybe renamed) and use that
-# to check if shardings are compatible with the input.
+# TODO(yashkatariya): Generalize check_compatible_aval (maybe renamed) and use
+# that to check if shardings are compatible with the input.
 @lru_cache(maxsize=2048)
 def _check_sharding(aval, s):
   if isinstance(s, Sharding):
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index b6283251e..f3755cb85 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -1294,10 +1294,10 @@ def pjit_check_aval_sharding(
     name_str = f' with pytree key path {name}' if name else ''
     shape = aval.shape
     try:
-      # Sharding interfaces can implement `is_compatible_aval` as an optional
+      # Sharding interfaces can implement `check_compatible_aval` as an optional
       # method to raise a more meaningful error.
-      if hasattr(s, 'is_compatible_aval'):
-        s.is_compatible_aval(shape)
+      if hasattr(s, 'check_compatible_aval'):
+        s.check_compatible_aval(shape)
       else:
         s._to_xla_hlo_sharding(len(shape))
     except ValueError as e:
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index c2860ed47..7a03e5ccc 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -322,7 +322,7 @@ class NamedSharding(XLACompatibleSharding):
       return False
     return self.mesh is other.mesh or self.mesh == other.mesh
 
-  def is_compatible_aval(self, aval_shape: Shape):
+  def check_compatible_aval(self, aval_shape: Shape) -> None:
     assert self._parsed_pspec is not None
     if len(aval_shape) < len(self._parsed_pspec):
       extra_msg = (' For scalars the PartitionSpec should be P()'
@@ -716,7 +716,7 @@ class PositionalSharding(XLACompatibleSharding):
     new_ids = self._ids.sum(axis=axis, keepdims=keepdims)  # union
     return self._remake(self._devices, new_ids)
 
-  def is_compatible_aval(self, aval_shape: Shape) -> None:
+  def check_compatible_aval(self, aval_shape: Shape) -> None:
     if len(aval_shape) != len(self.shape):
       raise ValueError(
           f"Sharding {self} is only valid for values of rank "
@@ -872,7 +872,7 @@ class GSPMDSharding(XLACompatibleSharding):
     mem = '' if self._memory_kind is None else f', memory_kind={self._memory_kind}'
     return f'GSPMDSharding({self._hlo_sharding!r}{mem})'
 
-  def is_compatible_aval(self, aval_shape: Shape):
+  def check_compatible_aval(self, aval_shape: Shape) -> None:
     num_ways_dim_sharded, _ = get_num_ways_dim_sharded(self._hlo_sharding)
     if len(aval_shape) < len(num_ways_dim_sharded):
       raise ValueError(


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index 5b975f58f..b2e2100d4 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -923,7 +923,7 @@ class ShardingTest(jtu.JaxTestCase):
         r"Sharding NamedSharding\(mesh=Mesh\('replica': 1, 'data': 1, 'mdl': 2\), "
         r"spec=PartitionSpec\(None, \('mdl',\), None, None\).*\) is only "
         "valid for values of rank at least 4, but was applied to a value of rank 2"):
-      new_mps.is_compatible_aval(shape)
+      new_mps.check_compatible_aval(shape)
 
   def test_is_subclass(self):
     # array version of api_test.py::APITest::test_is_subclass
@@ -1188,7 +1188,7 @@ class ShardingTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(
         ValueError,
         r"For scalars the PartitionSpec should be P()"):
-      s.is_compatible_aval(shape)
+      s.check_compatible_aval(shape)
 
   def test_mesh_caching_during_construction(self):
     if jax.device_count() < 2: