jax-ml/jax · error · NotImplementedError

mixes JAX mesh and Pallas mesh grid axes

Error message

{axis_name} mixes JAX mesh and Pallas mesh grid axes

What it means

In _device_id_dict_to_mesh, a joint (tuple) axis name in the device_id dict may not partially overlap the Pallas mesh axes: either every component axis is a mesh axis or none are. A tuple mixing JAX mesh axes and non-mesh (Pallas grid) axes is ambiguous and rejected.

Solutions

Split the joint axis so each tuple contains only JAX mesh axes or only Pallas grid axes
Use separate collectives per axis kind

Example fix

// before
axis_name=('mesh_axis', 'pallas_grid_axis')
// after
# perform collectives on ('mesh_axis','mesh_axis2') and 'pallas_grid_axis' separately

Defensive patterns

Strategy: validation

Validate before calling

for ax in axis_names:
    if isinstance(ax, tuple):
        kinds = {is_mesh_axis(a) for a in ax}  # your axis bookkeeping
        assert len(kinds) == 1, f"joint axis {ax} mixes JAX mesh and Pallas grid axes"

Prevention

Keep joint axis tuples homogeneous (all mesh or all grid)
Document which axis names belong to JAX Mesh vs Pallas grid

When it happens

Trigger: Declaring a collective whose axis name is a tuple containing both a jax.lax mesh axis and a Pallas grid axis name in the same tuple.

Common situations: Writing multi-axis collectives (e.g., ('data','rep') where 'data' is a JAX Mesh axis and 'rep' a Pallas grid axis) in distributed Pallas kernels.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/cb8b46c42c8ddb52. Report an issue: GitHub.

Appendix: source

Thrown at jax/_src/pallas/primitives.py:1215

    _semaphore_wait_discharge_rule
)


def _device_id_dict_to_mesh(mesh_context: pallas_utils.MeshInfo | None, device_id_dict, get_axis_index):
  if mesh_context is None:
    mesh_axis_sizes = {}
  else:
    mesh_axis_sizes = dict(
        zip(mesh_context.axis_names, mesh_context.mesh_shape)
    )
  physical_axis_dict = {}
  # Handle joint axes (i.e., one logical axis over >1 physical axes)
  for axis_name, idx in device_id_dict.items():
    if isinstance(axis_name, tuple) and any(
        a in mesh_axis_sizes for a in axis_name
    ):
      if not all(a in mesh_axis_sizes for a in axis_name):
        raise NotImplementedError(
            f"{axis_name} mixes JAX mesh and Pallas mesh grid axes"
        )
      axes_dimensions = [mesh_axis_sizes[name] for name in axis_name]
      for axis_index, axis_name in enumerate(axis_name):
        axis_size = mesh_axis_sizes[axis_name]
        inner_mesh_size = math.prod(axes_dimensions[axis_index + 1 :])

        # Fast path for power of 2s
        if inner_mesh_size & (inner_mesh_size - 1) == 0:
          shift_len = (inner_mesh_size & -inner_mesh_size).bit_length() - 1
          partial_device_idx = idx >> shift_len
        else:
          partial_device_idx = idx // inner_mesh_size

        if axis_size & (axis_size - 1) == 0:
          device_idx = partial_device_idx & jnp.asarray(
              axis_size - 1, dtype=partial_device_idx.dtype
          )

View on GitHub (pinned to 1e1c6a8fc0)