From b643e3994f259ed0c6c8789ee462fd21864d810f Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Fri, 23 Oct 2026 22:31:30 -0700
Subject: [PATCH] Validate per-axis divisibility in reshard.

PiperOrigin-RevId: 1006105626
---
 pathwaysutils/reshard.py           |  6 +++---
 pathwaysutils/test/reshard_test.py | 13 +++++++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/pathwaysutils/reshard.py b/pathwaysutils/reshard.py
index 64bf117..87f2b8a 100644
--- a/pathwaysutils/reshard.py
+++ b/pathwaysutils/reshard.py
@@ -58,9 +58,8 @@ def reshard_generic(
   for index, (arr, dst_sharding) in enumerate(zip(flat_x, flat_sharding)):
     if not isinstance(dst_sharding, jax.sharding.Sharding):
       raise ValueError("`sharding` must contain only `jax.sharding.Sharding`")
-    if not isinstance(arr, jax.Array) or (
-        hasattr(arr, "dtype")
-        and jax.dtypes.issubdtype(arr.dtype, jax.dtypes.prng_key)
+    if not isinstance(arr, jax.Array) or jax.dtypes.issubdtype(
+        arr.dtype, jax.dtypes.prng_key
     ):
       non_reshardable_arrays["arrays"].append(arr)
       non_reshardable_arrays["indices"].append(index)
@@ -212,6 +211,7 @@ def _get_split_candidates(
     if (
         gcd_shards[i] == 1
         and src_dims[i] > dst_dims[i]
+        and src_dims[i] % dst_dims[i] == 0
         and isinstance(spec, str)
     ):
       split_candidates.append((i, spec))
diff --git a/pathwaysutils/test/reshard_test.py b/pathwaysutils/test/reshard_test.py
index 0cc6c3f..b12b068 100644
--- a/pathwaysutils/test/reshard_test.py
+++ b/pathwaysutils/test/reshard_test.py
@@ -90,5 +90,18 @@ def test_ifrt_reshard_pytree(self):
     self.assertLen(args[0], 2)
     self.assertLen(args[1], 2)
 
+  def test_find_intermediate_sharding_rejects_indivisible_candidate(self):
+    # src_dims=[8, 3], dst_dims=[3, 4], gcd_shards=[1, 1] -> src_dims[0] % dst_dims[0] != 0
+    in_sharding = mock.MagicMock(spec=jax.sharding.NamedSharding)
+    in_sharding.spec = jax.sharding.PartitionSpec("x", "y")
+    with self.assertRaises(reshard.NoIntermediateShardingError):
+      reshard._get_split_candidates(
+          in_sharding,
+          src_dims=[8, 3],
+          dst_dims=[3, 4],
+          gcd_shards=[1, 1],
+      )
+
+
 if __name__ == "__main__":
   absltest.main()