jax-ml/jax · error · ValueError

Axes mentioned in `manual_axis_type` field of ShapedArray sh

Error message

Axes mentioned in `manual_axis_type` field of ShapedArray should be of type `Manual`. Got manual_axis_type={mat} with axis: {i} of type {mesh._name_to_type[i]}

What it means

ShapedArray's manual_axis_type (mat) lists mesh axes that are 'manual' (unsharded from this array's perspective). JAX validates that every axis named in the manual sets is actually declared with AxisType.Manual on the mesh; naming an Implicit/Explicit axis is a programming error.

Source

Thrown at jax/_src/core.py:2323

        "sharding with memory_kind is not allowed. Please use `jax.device_put`"
        f" to transfer to different memory spaces. Got {sharding=}")
  return out_s


@cache(max_size=4096,
       trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value)
def get_mat(mat, mesh):
  if mesh.empty:
    assert mat.empty, mat
    return mat

  axis_env = get_axis_env()
  in_axis_env = lambda i: axis_env.axis_exists(i) and i not in mesh._name_to_type
  for i in it.chain(mat.varying, mat.unreduced, mat.reduced):
    if in_axis_env(i):
      continue
    if mesh._name_to_type[i] != AxisType.Manual:
      raise ValueError(
          "Axes mentioned in `manual_axis_type` field of ShapedArray should be"
          f" of type `Manual`. Got manual_axis_type={mat} with axis: {i} of"
          f" type {mesh._name_to_type[i]}")
  if config.remove_size_one_mesh_axis_from_type.value:
    varying = frozenset(i for i in mat.varying
                        if in_axis_env(i) or mesh.shape[i] != 1)
    unreduced = frozenset(u for u in mat.unreduced if mesh.shape[u] != 1)
    reduced = frozenset(r for r in mat.reduced if mesh.shape[r] != 1)
    u_kind = mat.unreduced_kind if unreduced else None
    return mat.update(varying=varying, unreduced=unreduced, reduced=reduced,
                      unreduced_kind=u_kind)
  return mat


def get_memory_space(memory_space):
  assert memory_space is not None
  return memory_space

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Declare the axis as manual when creating the mesh: jax.make_mesh(..., axis_types=(AxisType.Manual, ...))
2. Remove the offending axis from manual_axis_type sets so it isn't claimed manual
3. Rebuild the mesh with correct axis_types before constructing avals

Example fix

// before
mesh = jax.make_mesh((4,), ('x',))  # not manual
mat = ManualAxisType(varying={'x'})  # rejected

// after
from jax.sharding import AxisType
mesh = jax.make_mesh((4,), ('x',), axis_types=(AxisType.Manual,))
mat = ManualAxisType(varying={'x'})

def axes_all_manual(mat, mesh):
    return all(mesh._name_to_type.get(i) == AxisType.Manual
               for i in (mat.varying | mat.unreduced | mat.reduced))

def valid_mat(mat, mesh):
    return not ((mat.varying | mat.unreduced | mat.reduced) - set(mesh._name_to_type)) and all(mesh._name_to_type[i] == AxisType.Manual for i in mat.varying|mat.unreduced|mat.reduced)

When it happens

Trigger: Constructing ShapedArray with a ManualAxisType containing an axis whose mesh._name_to_type maps to AxisType.Explicit or Implicit; manually building avals under spmd/shard_map code paths.

Common situations: Custom primitives or debug code fabricating avals; mesh created without axis_types (all axes default non-manual) then reused with manual mat annotations; refactors of AxisType APIs across JAX versions.

Related errors

• varying and unreduced cannot have common mesh axes. Got vary
• varying and reduced cannot have common mesh axes. Got varyin
• {name} only accepts inputs that are unreduced. Got {aval.str
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

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