{"record":{"id":"a02fd6ddc7de4114","repo":"jax-ml/jax","slug":"when-check-vma-true-on-jax-shard-map-manual","errorCode":null,"errorMessage":"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`.","messagePattern":"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`\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":1681,"sourceCode":"  with pallas_export_experimental(dynamic_shapes):\n    f = jit(f, device=device, static_argnames=static_argnames)\n    if platforms is None:\n      platforms = [\"tpu\"]\n    exported = export(f, platforms=platforms)(*args, **kwargs)\n    return exported.mlir_module()\n\n\n_out_shape_to_aval_mapping: dict[\n    type[Any], Callable[[Any], jax_core.AbstractValue]\n] = {}\n\n\ndef _convert_out_shape_to_aval(out_shape: Any) -> jax_core.AbstractValue:\n  match out_shape:\n    case jax_core.ShapeDtypeStruct():\n      if config._check_vma.value:\n        if out_shape.manual_axis_type is None:\n          raise ValueError(\n              \"When `check_vma=True` on `jax.shard_map`, `manual_axis_type` on\"\n              \" `jax.ShapeDtypeStruct` must not be `None`. Please specify how\"\n              \" the output should be varying across mesh axes using the\"\n              \" `manual_axis_type` argument of `jax.ShapeDtypeStruct` or set\"\n              \" `check_vma=False` on `jax.shard_map`.\")\n        return jax_core.ShapedArray(\n            shape=out_shape.shape, dtype=out_shape.dtype,\n            sharding=jax_core.get_cur_mesh_sharding(),\n            manual_axis_type=out_shape.manual_axis_type)\n      return jax_core.ShapedArray(\n          shape=out_shape.shape, dtype=out_shape.dtype,\n          sharding=jax_core.get_cur_mesh_sharding())\n    case jax_core.ShapedArray():\n      return out_shape\n    case MemoryRef():\n      return out_shape.get_array_aval()\n    case hijax.HiType():\n      return out_shape","sourceCodeStart":1663,"sourceCodeEnd":1699,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L1663-L1699","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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","If you cannot determine axis types yet, set check_vma=False on jax.shard_map to skip validation","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","Audit each mesh axis and annotate outputs: varying (replicated logic per shard) vs manual (per-shard slices) to keep validation on"],"exampleFix":"# before\nout = jax.ShapeDtypeStruct((8, 8), jnp.float32)\nf_smap = jax.shard_map(f, mesh, out_shape=out)  # check_vma on -> ValueError\n\n# after\nfrom jax.sharding import AxisType\nout = jax.ShapeDtypeStruct((8, 8), jnp.float32, manual_axis_type=(AxisType.Varying, AxisType.Varying))\nf_smap = jax.shard_map(f, mesh, out_shape=out)","handlingStrategy":"validation","validationCode":"import jax\nout = jax.ShapeDtypeStruct(shape, dtype)\nif jax.config._check_vma.value and out.manual_axis_type is None:\n    raise ValueError('set manual_axis_type or pass check_vma=False')  # fail fast with your own message","typeGuard":"def out_shape_is_vma_complete(out_shape) -> bool:\n    return out_shape.manual_axis_type is not None and len(out_shape.manual_axis_type) == len(out_shape.shape)","tryCatchPattern":"try:\n    f_smap = jax.shard_map(f, mesh, out_shape=out)\nexcept ValueError as e:\n    if 'manual_axis_type' in str(e):\n        f_smap = jax.shard_map(f, mesh, out_shape=out, check_vma=False)\n    else:\n        raise","preventionTips":["Always annotate manual_axis_type on out_shapes in VMA-enabled codebases","Centralize check_vma policy in one config so all shard_map calls agree","Add a unit test asserting out_shapes are VMA-complete under jax_check_vma=True"],"tags":["jax","shard-map","vma","shapedtypestruct","mesh","tpu","config-flag"],"backgroundTag":"missing-required-annotation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}