jax-ml/jax · error · NotImplementedError

{axis_name} 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_to_logical, a joint axis name (tuple of axes) must be composed entirely of JAX mesh axes or entirely of Pallas mesh grid axes. Mixing both kinds inside one joint axis is not implemented.

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)

Solutions

1. Split the joint axis so each tuple contains only JAX mesh axes or only Pallas grid axes
2. 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

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"

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

• unbound axis name: {axis_name}
• all_gather_reduced is a Varying -> Reduced collective. This
• unreduced_psum_scatter is a Unreduced -> Varying collective.
• {name} is a Unreduced -> Invariant collective. This means th
• run_scoped_p with collective axes is not supported

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