custom_parititioning: in lower sharding, Sharding should be XLACompat…

…ibleSharding. PiperOrigin-RevId: 537077304
jakevdp · Jun 1, 2023 · 5c2070c · 5c2070c
1 parent c8311c6
commit 5c2070c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/jax/experimental/custom_partitioning.py b/jax/experimental/custom_partitioning.py
@@ -121,8 +121,10 @@ def _custom_partitioning_propagate_user_sharding(user_sharding, shape,
 
 
 def _to_hlo_sharding(sharding, num_dimensions):
-  if not isinstance(sharding, jax.sharding.Sharding):
-    raise ValueError("Custom Partitioning rules must return shardings.")
+  if not isinstance(sharding, jax.sharding.XLACompatibleSharding):
+    raise ValueError(
+        "Custom Partitioning rules must return XLACompatibleShardings."
+    )
   return xc.HloSharding.from_proto(sharding._to_xla_op_sharding(num_dimensions))