Fixed shard ordering mismatch during weights sync.

tpu_raiden/api/jax/BUILD

@@ -75,10 +75,15 @@ py_test(
 
 py_library(
     name = "weight_synchronizer_jax_py",
-    srcs = ["weight_synchronizer.py"],
+    srcs = [
+        "utils.py",
+        "weight_synchronizer.py",
+    ],
     visibility = ["//visibility:public"],
     deps = [
         "//tpu_raiden/frameworks/jax:_tpu_raiden_jax",
+        "@jax//jax",
+        "@pypi//numpy",
     ],
 )
 

tpu_raiden/api/jax/utils.py

@@ -0,0 +1,149 @@
+# Copyright 2026 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
+#
+#     http://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.
+
+# Copyright 2026 The TPU Raiden Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+"""Utility functions for JAX Weight Synchronizer."""
+
+from typing import List
+import jax
+import numpy as np
+
+
+def get_shard_sorting_permutation(arr: jax.Array) -> List[int]:
+  """Computes the permutation to sort JAX shards for Raiden Controller."""
+  sharding = arr.sharding
+  if not isinstance(sharding, jax.sharding.NamedSharding):
+    return list(range(len(arr.addressable_shards)))
+
+  mesh = sharding.mesh
+  spec = sharding.spec
+
+  # 1. Reconstruct logical_mesh_shape and layout
+  logical_mesh_shape = [mesh.shape[ax] for ax in mesh.axis_names]
+
+  sharded_axes = []
+  for axis in spec:
+    if axis is None:
+      continue
+    if isinstance(axis, str):
+      sharded_axes.append(axis)
+    else:
+      sharded_axes.extend(axis)
+
+  major_to_minor = sorted([mesh.axis_names.index(ax) for ax in sharded_axes])
+
+  # 2. Compute controller's expected global indices
+  num_shards = len(arr.addressable_shards)
+  num_physical_hosts = jax.process_count()
+  replica_id = jax.process_index()
+
+  phys_mesh = [logical_mesh_shape[d] for d in major_to_minor]
+
+  host_axis_logical = None
+  for d, size in enumerate(logical_mesh_shape):
+    if size == num_physical_hosts:
+      host_axis_logical = d
+      break
+
+  non_host_axes = [
+      d for d in range(len(logical_mesh_shape)) if d != host_axis_logical
+  ]
+
+  controller_global_indices = []
+  for j in range(num_shards):
+    local_coords = {}
+    temp = j
+    for d in reversed(non_host_axes):
+      size = logical_mesh_shape[d]
+      local_coords[d] = temp % size
+      temp = temp // size
+
+    full_coords = [0] * len(logical_mesh_shape)
+    for d in range(len(logical_mesh_shape)):
+      if d == host_axis_logical:
+        full_coords[d] = replica_id
+      else:
+        full_coords[d] = local_coords.get(d, 0)
+
+    tensor_coords = [full_coords[m_axis] for m_axis in major_to_minor]
+
+    global_idx = 0
+    stride = 1
+    for val, size in zip(reversed(tensor_coords), reversed(phys_mesh)):
+      global_idx += val * stride
+      stride *= size
+    controller_global_indices.append(global_idx)
+
+  # 3. Compute ACTUAL JAX global shard indices
+  jax_shard_global_indices = []
+  for shard in arr.addressable_shards:
+    device = shard.device
+    coords = np.argwhere(mesh.devices == device)
+    if coords.size == 0:
+      raise ValueError(f"Device {device} not found in mesh")
+    m_coords = coords[0]
+    full_coords = list(m_coords)
+
+    # Map array dimensions to mesh axes using spec
+    tensor_coords = []
+    tensor_shape = []
+    for axis in spec:
+      if axis is None:
+        tensor_coords.append(0)
+        tensor_shape.append(1)
+      elif isinstance(axis, str):
+        tensor_coords.append(full_coords[mesh.axis_names.index(axis)])
+        tensor_shape.append(mesh.shape[axis])
+      else:
+        for ax in axis:
+          tensor_coords.append(full_coords[mesh.axis_names.index(ax)])
+          tensor_shape.append(mesh.shape[ax])
+
+    # Compute flat index (row-major)
+    global_idx = 0
+    stride = 1
+    for val, size in zip(reversed(tensor_coords), reversed(tensor_shape)):
+      global_idx += val * stride
+      stride *= size
+    jax_shard_global_indices.append(global_idx)
+
+  # 4. Sort indices
+  indices = list(range(len(arr.addressable_shards)))
+
+  def sort_key(idx):
+    g = jax_shard_global_indices[idx]
+    occurrence = jax_shard_global_indices[:idx].count(g)
+    matching_indices = [
+        i for i, val in enumerate(controller_global_indices) if val == g
+    ]
+    if occurrence < len(matching_indices):
+      target_j = matching_indices[occurrence]
+    else:
+      target_j = matching_indices[-1]
+    return target_j
+
+  sorted_indices = sorted(indices, key=sort_key)
+  return sorted_indices

tpu_raiden/api/jax/weight_synchronizer_test.py

@@ -22,6 +22,7 @@
 import jax.numpy as jnp
 import numpy as np
 
+from tpu_raiden.api.jax import utils
 from tpu_raiden.api.jax import weight_synchronizer
 from tpu_raiden.rpc import raiden_service_pb2
 
@@ -43,6 +44,9 @@ def setUp(self):
     self.sharding = jax.sharding.NamedSharding(
         self.mesh, jax.sharding.PartitionSpec("data")
     )
+    self.mesh_2d = jax.sharding.Mesh(
+        np.array(self.devices[:4]).reshape(2, 2), ("x", "y")
+    )
     self.shape = (8, 128)
     self.dtype = jnp.float32
 
@@ -216,6 +220,110 @@ def test_bind_weights(self):
     for arr in dst_arrs:
       np.testing.assert_array_equal(np.asarray(arr), 5.0)
 
+  def _run_resharding_test(self, src_sharding, dst_sharding, shape):
+    src_arrs = [
+        jax.device_put(
+            jnp.arange(np.prod(shape), dtype=self.dtype).reshape(shape),
+            src_sharding,
+        )
+    ]
+    dst_arrs = [
+        jax.device_put(jnp.zeros(shape, dtype=self.dtype), dst_sharding)
+    ]
+
+    for arr in src_arrs:
+      arr.block_until_ready()
+    for arr in dst_arrs:
+      arr.block_until_ready()
+
+    ws_source = WeightSynchronizer(
+        jax_arrays=src_arrs,
+        local_port=0,
+        unsafe_skip_buffer_lock=True,
+        listener_port=0,
+        bind_ip="127.0.0.1",
+    )
+    ws_dest = WeightSynchronizer(
+        jax_arrays=dst_arrs,
+        local_port=0,
+        unsafe_skip_buffer_lock=True,
+        bind_ip="127.0.0.1",
+    )
+
+    req = raiden_service_pb2.ControlRequest(
+        command=raiden_service_pb2.ControlRequest.COMMAND_START_TRANSFER,
+        peers=[
+            f"127.0.0.1:{ws_dest.local_port}",
+        ],
+        start_transfer_request=raiden_service_pb2.StartTransferRequest(
+            is_sender=True
+        ),
+    )
+    payload = req.SerializeToString()
+
+    sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM, 0)
+    sock.connect(("::1", ws_source.listener_port))
+    sock.sendall(len(payload).to_bytes(4, "big") + payload)
+
+    resp_len = int.from_bytes(sock.recv(4), "big")
+    resp_bytes = sock.recv(resp_len)
+    resp = raiden_service_pb2.ControlResponse()
+    resp.ParseFromString(resp_bytes)
+    self.assertTrue(resp.success)
+    sock.close()
+
+    ws_dest.h2d()
+
+    # Verify data integrity
+    np.testing.assert_array_equal(
+        np.asarray(dst_arrs[0]), np.asarray(src_arrs[0])
+    )
+
+  def test_push_sync_aligned_to_aligned(self):
+    src_sharding = jax.sharding.NamedSharding(
+        self.mesh_2d, jax.sharding.PartitionSpec("x", "y")
+    )
+    dst_sharding = jax.sharding.NamedSharding(
+        self.mesh_2d, jax.sharding.PartitionSpec("x", "y")
+    )
+    self._run_resharding_test(src_sharding, dst_sharding, (8, 8))
+
+class ShardSortingUtilTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    try:
+      self.devices = jax.devices("tpu")
+    except RuntimeError:
+      self.devices = jax.devices("cpu")
+    self.mesh_2d = jax.sharding.Mesh(
+        np.array(self.devices[:4]).reshape(2, 2), ("x", "y")
+    )
+
+  def test_aligned_sharding_permutation(self):
+    sharding = jax.sharding.NamedSharding(
+        self.mesh_2d, jax.sharding.PartitionSpec("x", "y")
+    )
+    arr = jax.device_put(jnp.zeros((8, 8)), sharding)
+    perm = utils.get_shard_sorting_permutation(arr)
+    self.assertEqual(perm, [0, 1, 2, 3])
+
+  def test_transposed_sharding_permutation(self):
+    sharding = jax.sharding.NamedSharding(
+        self.mesh_2d, jax.sharding.PartitionSpec("y", "x")
+    )
+    arr = jax.device_put(jnp.zeros((8, 8)), sharding)
+    perm = utils.get_shard_sorting_permutation(arr)
+    self.assertEqual(perm, [0, 2, 1, 3])
+
+  def test_replicated_sharding_permutation(self):
+    sharding = jax.sharding.NamedSharding(
+        self.mesh_2d, jax.sharding.PartitionSpec("x")
+    )
+    arr = jax.device_put(jnp.zeros((8, 8)), sharding)
+    perm = utils.get_shard_sorting_permutation(arr)
+    self.assertEqual(perm, [0, 1, 2, 3])
+
 
 if __name__ == "__main__":
   absltest.main()

tpu_raiden/frameworks/jax/utils.h

@@ -51,6 +51,29 @@ inline std::vector<std::vector<raiden::RaidenBufferHandle>> UnpackJaxArrays(
           "Number of shards mismatch across layers during unpack");
     }
 
+    try {
+      nanobind::object utils_mod = nanobind::module_::import_(
+          "tpu_raiden.api.jax.utils");
+      nanobind::object get_permutation_fn =
+          utils_mod.attr("get_shard_sorting_permutation");
+      nanobind::list permutation =
+          nanobind::cast<nanobind::list>(get_permutation_fn(dst));
+
+      std::vector<raiden::RaidenBufferHandle> sorted_shard_buffers;
+      sorted_shard_buffers.reserve(num_shards);
+      for (size_t i = 0; i < num_shards; ++i) {
+        size_t src_idx = nanobind::cast<size_t>(permutation[i]);
+        if (src_idx >= shard_buffers.size()) {
+          throw std::runtime_error("Permutation index out of bounds");
+        }
+        sorted_shard_buffers.push_back(std::move(shard_buffers[src_idx]));
+      }
+      shard_buffers = std::move(sorted_shard_buffers);
+    } catch (const std::exception& e) {
+      throw std::runtime_error(std::string("Failed to sort shards: ") +
+                               e.what());
+    }
+
     layer_buffers.push_back(std::move(shard_buffers));
   }
   return layer_buffers;