jax-ml/jax · error · ValueError

When `check_vma=True` on `jax.shard_map`, `manual_axis_type`

Error message

When `check_vma=True` on `jax.shard_map`, `manual_axis_type` on `jax.ShapeDtypeStruct` must not be `None`. Please specify how the output should be varying across mesh axes using the `manual_axis_type` argument of `jax.ShapeDtypeStruct` or set `check_vma=False` on `jax.shard_map`.

What it means

When the config flag jax_check_vma (check_vma) is enabled, jax.shard_map output shapes declared as ShapeDtypeStruct must specify manual_axis_type, which describes how each output axis varies across mesh axes. Without it, shard_map cannot verify that manually specified out_shapes agree with VMA (varying-manual-axis) rules, so it raises ValueError.

Source

Thrown at jax/_src/pallas/core.py:1681

  with pallas_export_experimental(dynamic_shapes):
    f = jit(f, device=device, static_argnames=static_argnames)
    if platforms is None:
      platforms = ["tpu"]
    exported = export(f, platforms=platforms)(*args, **kwargs)
    return exported.mlir_module()


_out_shape_to_aval_mapping: dict[
    type[Any], Callable[[Any], jax_core.AbstractValue]
] = {}


def _convert_out_shape_to_aval(out_shape: Any) -> jax_core.AbstractValue:
  match out_shape:
    case jax_core.ShapeDtypeStruct():
      if config._check_vma.value:
        if out_shape.manual_axis_type is None:
          raise ValueError(
              "When `check_vma=True` on `jax.shard_map`, `manual_axis_type` on"
              " `jax.ShapeDtypeStruct` must not be `None`. Please specify how"
              " the output should be varying across mesh axes using the"
              " `manual_axis_type` argument of `jax.ShapeDtypeStruct` or set"
              " `check_vma=False` on `jax.shard_map`.")
        return jax_core.ShapedArray(
            shape=out_shape.shape, dtype=out_shape.dtype,
            sharding=jax_core.get_cur_mesh_sharding(),
            manual_axis_type=out_shape.manual_axis_type)
      return jax_core.ShapedArray(
          shape=out_shape.shape, dtype=out_shape.dtype,
          sharding=jax_core.get_cur_mesh_sharding())
    case jax_core.ShapedArray():
      return out_shape
    case MemoryRef():
      return out_shape.get_array_aval()
    case hijax.HiType():
      return out_shape

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set manual_axis_type on the out_shape, e.g. ShapeDtypeStruct(shape, dtype, manual_axis_type=(AxisType.Varying/Manual per axis)) matching how the function output varies across mesh axes
2. If you cannot determine axis types yet, set check_vma=False on jax.shard_map to skip validation
3. Set the config flag off via jax.config.update('jax_check_vma', False) or drop JAX_CHECK_VMA=1 from the environment if the check is not required
4. Audit each mesh axis and annotate outputs: varying (replicated logic per shard) vs manual (per-shard slices) to keep validation on

Example fix

# before
out = jax.ShapeDtypeStruct((8, 8), jnp.float32)
f_smap = jax.shard_map(f, mesh, out_shape=out)  # check_vma on -> ValueError

# after
from jax.sharding import AxisType
out = jax.ShapeDtypeStruct((8, 8), jnp.float32, manual_axis_type=(AxisType.Varying, AxisType.Varying))
f_smap = jax.shard_map(f, mesh, out_shape=out)

import jax
out = jax.ShapeDtypeStruct(shape, dtype)
if jax.config._check_vma.value and out.manual_axis_type is None:
    raise ValueError('set manual_axis_type or pass check_vma=False')  # fail fast with your own message

def out_shape_is_vma_complete(out_shape) -> bool:
    return out_shape.manual_axis_type is not None and len(out_shape.manual_axis_type) == len(out_shape.shape)

try:
    f_smap = jax.shard_map(f, mesh, out_shape=out)
except ValueError as e:
    if 'manual_axis_type' in str(e):
        f_smap = jax.shard_map(f, mesh, out_shape=out, check_vma=False)
    else:
        raise

When it happens

Trigger: Calling jax.shard_map(f, mesh, out_shape=jax.ShapeDtypeStruct(shape, dtype)) with jax_check_vma=True (explicitly or via JAX_CHECK_VMA=1) and manual_axis_type left as None on the out_shape.

Common situations: Enabling VMA checking in multi-host/multi-DC jax experiments (common in TPU pod training repos); upgrading jax where check_vma defaults or shard_map signatures changed; copying old shard_map code that predates manual_axis_type.

Related errors

• Mesh must be provided for shard_map with checkify.
• {} function carry input and carry output must have equal typ
• Scan carry input and output got mismatched varying manual ax
• Found an unbound axis name: {name}. To fix this, please call
• When passing host local inputs to pjit, devices connected to

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