Refactor resharding logic into helper functions.

This change introduces `_reshard_with_sidechannel` and `_reshard_with_ifrt` to encapsulate the different resharding mechanisms used by the `reshard` function. These are internal APIs and should not be depended on.

PiperOrigin-RevId: 856749135

Author: lukebaumann

pathwaysutils/experimental/reshard.py

@@ -198,6 +198,42 @@ def _ifrt_jax_array_reshard(
   )
 
 
+def _reshard_with_sidechannel(
+    x: Any,
+    sharding: jax.sharding.Sharding | Any,
+    *,
+    donate: bool = False,
+    may_alias: bool | None = None,
+    cache_resharding_plans: bool = False,
+) -> Any:
+  """Reshards `x` to `sharding` using sidechannel."""
+  return _reshard(
+      x,
+      sharding,
+      donate=donate,
+      may_alias=may_alias,
+      jax_array_reshard_fn=_sidechannel_jax_array_reshard,
+      cache_resharding_plans=cache_resharding_plans,
+  )
+
+
+def _reshard_with_ifrt(
+    x: Any,
+    sharding: jax.sharding.Sharding | Any,
+    *,
+    donate: bool = False,
+    may_alias: bool | None = None,
+) -> Any:
+  """Reshards `x` to `sharding` using IFRT."""
+  return _reshard(
+      x,
+      sharding,
+      donate=donate,
+      may_alias=may_alias,
+      jax_array_reshard_fn=_ifrt_jax_array_reshard,
+  )
+
+
 def reshard(
     x: Any,
     sharding: jax.sharding.Sharding | Any,
@@ -231,19 +267,17 @@ def reshard(
     A copy of `x` whose sharding is `sharding`.
   """
   if pw_jax.ifrt_reshard_available():
-    return _reshard(
+    return _reshard_with_ifrt(
         x,
         sharding,
         donate=donate,
         may_alias=may_alias,
-        jax_array_reshard_fn=_ifrt_jax_array_reshard,
     )
   else:
-    return _reshard(
+    return _reshard_with_sidechannel(
         x,
         sharding,
         donate=donate,
         may_alias=may_alias,
-        jax_array_reshard_fn=_sidechannel_jax_array_reshard,
         cache_resharding_plans=cache_resharding_plans,
     )

pathwaysutils/test/experimental/reshard_test.py

@@ -0,0 +1,68 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the resharding utilities in `pathwaysutils.experimental.reshard`."""
+
+import json
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+import jax.numpy as jnp
+from pathwaysutils import jax as pw_jax
+from pathwaysutils import plugin_executable
+from pathwaysutils.experimental import reshard
+
+
+class ReshardTest(parameterized.TestCase):
+
+  @parameterized.parameters(True, False)
+  def test_reshard_plan_wrapper_donate(self, donate):
+    avals = [jax.core.ShapedArray((2, 2), jnp.float32)]
+    devices = jax.devices()
+    sharding = jax.sharding.SingleDeviceSharding(devices[0])
+
+    mock_pe = self.enter_context(
+        mock.patch.object(plugin_executable, 'PluginExecutable')
+    )
+
+    reshard.ReshardingPlanWrapper(avals, [sharding], [sharding], donate=donate)
+
+    mock_pe.assert_called_once()
+    args, _ = mock_pe.call_args
+    request = json.loads(args[0])
+    self.assertEqual(request['reshardRequest']['donateInput'], donate)
+
+  @parameterized.parameters(True, False)
+  def test_ifrt_reshard_donate(self, donate):
+    x = jnp.array([1, 2])
+    devices = jax.devices()
+    sharding = jax.sharding.SingleDeviceSharding(devices[0])
+
+    mock_transfer = self.enter_context(
+        mock.patch.object(pw_jax, 'transfer_to_shardings')
+    )
+    self.enter_context(
+        mock.patch.object(pw_jax, 'ifrt_reshard_available', return_value=True)
+    )
+
+    reshard.reshard(x, sharding, donate=donate)
+    # Check that transfer_to_shardings was called with donate=True
+    # Signature: transfer_to_shardings(arrays, shardings, donate)
+    mock_transfer.assert_called_with(mock.ANY, mock.ANY, donate)
+
+
+if __name__ == '__main__':
+  absltest.main()