diff --git a/pathwaysutils/experimental/concatenate_by_mesh_axis.py b/pathwaysutils/experimental/concatenate_by_mesh_axis.py
index e2a5804..0fd951e 100644
--- a/pathwaysutils/experimental/concatenate_by_mesh_axis.py
+++ b/pathwaysutils/experimental/concatenate_by_mesh_axis.py
@@ -68,10 +68,21 @@ def _get_named_sharding(array: jax.Array) -> jax.sharding.NamedSharding:
   if not flattened_arrays[0]:
     return array_trees[0]
 
-  # Extract the shared mesh from each PyTree (from an arbitrary array in each).
+  # Extract the shared mesh from each PyTree and verify all leaves in a PyTree
+  # share that mesh.
   meshes_to_concatenate = [
       _get_named_sharding(array).mesh for array in input_flat_arrays[0]
   ]
+  for tree_arrays, expected_mesh in zip(
+      flattened_arrays, meshes_to_concatenate
+  ):
+    for array in tree_arrays:
+      sharding = _get_named_sharding(array)
+      if sharding.mesh != expected_mesh:
+        raise ValueError(
+            "All arrays in each PyTree must have the same mesh. Got"
+            f" {sharding.mesh} vs. {expected_mesh}."
+        )
 
   # Validate that the meshes are compatible.
   reference_mesh = meshes_to_concatenate[0]
@@ -117,6 +128,19 @@ def _get_output_sharding(
   ) -> jax.sharding.NamedSharding:
     reference_sharding = _get_named_sharding(arrays[0])
     reference_spec = reference_sharding.spec
+    for array in arrays[1:]:
+      sharding = _get_named_sharding(array)
+      if sharding.spec != reference_spec:
+        raise ValueError(
+            "Corresponding arrays across PyTrees must have the same"
+            f" PartitionSpec. Got {sharding.spec} vs. {reference_spec}."
+        )
+      if sharding.memory_kind != reference_sharding.memory_kind:
+        raise ValueError(
+            "Corresponding arrays across PyTrees must have the same"
+            f" memory_kind. Got {sharding.memory_kind} vs."
+            f" {reference_sharding.memory_kind}."
+        )
     return jax.sharding.NamedSharding(
         concatenated_mesh,
         reference_spec,
diff --git a/pathwaysutils/experimental/split_by_mesh_axis.py b/pathwaysutils/experimental/split_by_mesh_axis.py
index 14b7b8a..79a21a5 100644
--- a/pathwaysutils/experimental/split_by_mesh_axis.py
+++ b/pathwaysutils/experimental/split_by_mesh_axis.py
@@ -129,7 +129,7 @@ def split_by_mesh_axis(
         range(axis_size, mesh.axis_sizes[mesh_axis_idx] + 1, axis_size)
     )
   else:
-    mesh_axis_sections = mesh_axis_indices_or_sections
+    mesh_axis_sections = list(mesh_axis_indices_or_sections)
     for i, boundary in enumerate(mesh_axis_sections):
       if boundary <= 0 or boundary >= mesh.axis_sizes[mesh_axis_idx]:
         raise ValueError(
@@ -142,7 +142,7 @@ def split_by_mesh_axis(
             "Mesh axis sections must be monotonically increasing, but got"
             f" {mesh_axis_sections=}."
         )
-    mesh_axis_sections += [mesh.axis_sizes[mesh_axis_idx]]  # pyrefly: ignore[unsupported-operation]
+    mesh_axis_sections += [mesh.axis_sizes[mesh_axis_idx]]
 
   submeshes = []
   axis_boundary_start = 0