{"record":{"id":"83732e3744045b00","repo":"jax-ml/jax","slug":"axes-mentioned-in-manual-axis-type-field-of-shap","errorCode":null,"errorMessage":"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]}","messagePattern":"Axes mentioned in `manual_axis_type` field of ShapedArray should be of type `Manual`\\. Got manual_axis_type=(.+?) with axis: (.+?) of type (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2323,"sourceCode":"        \"sharding with memory_kind is not allowed. Please use `jax.device_put`\"\n        f\" to transfer to different memory spaces. Got {sharding=}\")\n  return out_s\n\n\n@cache(max_size=4096,\n       trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value)\ndef get_mat(mat, mesh):\n  if mesh.empty:\n    assert mat.empty, mat\n    return mat\n\n  axis_env = get_axis_env()\n  in_axis_env = lambda i: axis_env.axis_exists(i) and i not in mesh._name_to_type\n  for i in it.chain(mat.varying, mat.unreduced, mat.reduced):\n    if in_axis_env(i):\n      continue\n    if mesh._name_to_type[i] != AxisType.Manual:\n      raise ValueError(\n          \"Axes mentioned in `manual_axis_type` field of ShapedArray should be\"\n          f\" of type `Manual`. Got manual_axis_type={mat} with axis: {i} of\"\n          f\" type {mesh._name_to_type[i]}\")\n  if config.remove_size_one_mesh_axis_from_type.value:\n    varying = frozenset(i for i in mat.varying\n                        if in_axis_env(i) or mesh.shape[i] != 1)\n    unreduced = frozenset(u for u in mat.unreduced if mesh.shape[u] != 1)\n    reduced = frozenset(r for r in mat.reduced if mesh.shape[r] != 1)\n    u_kind = mat.unreduced_kind if unreduced else None\n    return mat.update(varying=varying, unreduced=unreduced, reduced=reduced,\n                      unreduced_kind=u_kind)\n  return mat\n\n\ndef get_memory_space(memory_space):\n  assert memory_space is not None\n  return memory_space\n","sourceCodeStart":2305,"sourceCodeEnd":2341,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2305-L2341","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Declare the axis as manual when creating the mesh: jax.make_mesh(..., axis_types=(AxisType.Manual, ...))","Remove the offending axis from manual_axis_type sets so it isn't claimed manual","Rebuild the mesh with correct axis_types before constructing avals"],"exampleFix":"// before\nmesh = jax.make_mesh((4,), ('x',))  # not manual\nmat = ManualAxisType(varying={'x'})  # rejected\n\n// after\nfrom jax.sharding import AxisType\nmesh = jax.make_mesh((4,), ('x',), axis_types=(AxisType.Manual,))\nmat = ManualAxisType(varying={'x'})","handlingStrategy":"type-guard","validationCode":"def axes_all_manual(mat, mesh):\n    return all(mesh._name_to_type.get(i) == AxisType.Manual\n               for i in (mat.varying | mat.unreduced | mat.reduced))","typeGuard":"def valid_mat(mat, mesh):\n    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)","tryCatchPattern":null,"preventionTips":["Create meshes with explicit axis_types upfront","Don't hand-construct ManualAxisType unless you own the mesh definition"],"tags":["jax","sharding","manual-axis-type","mesh","validation"],"backgroundTag":"invalid-sharding-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}