jax-ml/jax · error · ValueError

pallas_call requires all mesh axes to be Manual, got {get_ab

Error message

pallas_call requires all mesh axes to be Manual, got {get_abstract_mesh().axis_types}

What it means

pallas_call operates on manual (per-device, unsharded) data. During abstract evaluation it checks that no ShapedArray input or output has a non-None sharding spec entry (i.e. no named/auto sharding along mesh axes); if any value is sharded across an axis, it raises ValueError showing the abstract mesh's axis types. The note in source acknowledges this check does not catch auto-mode non-manual axes at this point.

Source

Thrown at jax/_src/pallas/pallas_call.py:135

  # Make sure we don't return ShapedArray with pallas memory space to the
  # outside world.
  out_avals = tuple(a.update(memory_space=jax_core.MemorySpace.Device)
                    if isinstance(a, jax_core.ShapedArray) else a
                    for a in out_avals)

  # TODO(mattjj,yashkatariya): if we hide vmapped away mesh axes, use this:
  # if not (all(a.sharding.mesh.are_all_axes_manual for a in avals) and
  #         all(a.sharding.mesh.are_all_axes_manual for a in out_avals) and
  #         get_abstract_mesh().are_all_axes_manual):
  #   raise ValueError("pallas_call requires all mesh axes to be Manual, "
  #                    f"got {get_abstract_mesh().axis_types}")

  # NOTE(mattjj,yashkatariya): this doesn't catch auto-mode non-manual axes
  if not (all(p is None for a in avals if isinstance(a, jax_core.ShapedArray)
              for p in a.sharding.spec) and
          all(p is None for a in out_avals if isinstance(a, jax_core.ShapedArray)
              for p in a.sharding.spec)):
    raise ValueError("pallas_call requires all mesh axes to be Manual, "
                     f"got {get_abstract_mesh().axis_types}")
  return out_avals, effs


pallas_call_p.def_effectful_abstract_eval(_pallas_call_abstract_eval)

def _pallas_call_is_high(*_, jaxpr, **params):
  del params
  return jaxpr.is_high
pallas_call_p.is_high = _pallas_call_is_high


def _get_index_mapping(avals) -> dict[int, tuple[int, ...]]:
  indices = {}
  counter = 0
  for i, in_aval in enumerate(avals):
    local_counter = []
    for _ in range(len(in_aval.lo_ty())):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the pallas_call invocation in a manual-replicated/manual mesh scope (e.g. with mesh(manual_axes...): or jax.lax.with_sharding_constraint-free manual context) so all axes are Manual
2. Remove NamedSharding/PartitionSpec annotations from the arrays and out_shape passed to pallas_call
3. Run pallas_call outside the sharded/jit-with-mesh context, operating on fully replicated per-device buffers

Example fix

# before
with mesh(mesh_obj, ('x',)):  # Named axis
    pallas_call(kernel, grid=grid, out_shape=out_shape)(x)
# after
with mesh(mesh_obj):  # all axes manual
    pallas_call(kernel, grid=grid, out_shape=out_shape)(x)

from jax.sharding import NamedSharding

def all_manual(vals):
    for v in vals:
        s = getattr(v, 'sharding', None)
        if s is not None and any(p is not None for p in s.spec):
            return False
    return True
assert all_manual(list(args) + list(out_shapes))

def is_unsharded(x) -> bool:
    s = getattr(x, 'sharding', None)
    return s is None or all(p is None for p in s.spec)

When it happens

Trigger: Calling pallas_call with arguments or out_shape values that carry a NamedSharding/PartitionSpec placing data on mesh axes, while the abstract mesh has Named or Auto axes (not fully Manual).

Common situations: Using pallas_call inside jax.jit with a mesh context (e.g. via mesh() context manager or sharding annotations) intended for sharded SPMD code; forgetting to enter a manual mesh scope before calling Pallas kernels; mixing new sharding APIs with Pallas.

Related errors

• callbacks are only supported in spmd computations when all m
• pallas_call with a mesh does not support batching
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• callbacks do not support specifying sharding inside spmd com

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