{"record":{"id":"e10d6c9089a4cdf0","repo":"jax-ml/jax","slug":"axis-types-passed-to-name-must-be-of-type-jax-s","errorCode":null,"errorMessage":"axis_types passed to {name} must be of type `jax.sharding.AxisType`. Got {axis_types} of type {tuple(type(a) for a in axis_types)}","messagePattern":"axis_types passed to (.+?) must be of type `jax\\.sharding\\.AxisType`\\. Got (.+?) of type (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/mesh.py","lineNumber":126,"sourceCode":"  return Mesh(global_mesh.devices[subcube_indices_tuple], global_mesh.axis_names)\n\n\nclass AxisType(enum.Enum):\n  Auto = enum.auto()\n  Explicit = enum.auto()\n  Manual = enum.auto()\n\n  def __repr__(self):\n    return self.name\n\ndef _normalize_axis_types(axis_names, axis_types, name, default_axis_type):\n  axis_types = ((default_axis_type,) * len(axis_names)\n                if axis_types is None else axis_types)\n  if not isinstance(axis_types, tuple):\n    axis_types = (axis_types,)\n\n  if not all(isinstance(a, AxisType) for a in axis_types):\n    raise TypeError(\n        f\"axis_types passed to {name} must be of type `jax.sharding.AxisType`.\"\n        f\" Got {axis_types} of type {tuple(type(a) for a in axis_types)}\")\n  if len(axis_names) != len(axis_types):\n    raise ValueError(\n        \"Number of axis names should match the number of axis_types. Got\"\n        f\" axis_names={axis_names} and axis_types={axis_types}\")\n  return axis_types\n\ndef all_axis_types_match(axis_types, ty: AxisType) -> bool:\n  if not axis_types:\n    return False\n  return all(t == ty for t in axis_types)\n\ndef any_axis_types_match(axis_types, ty: AxisType) -> bool:\n  if not axis_types:\n    return False\n  return any(t == ty for t in axis_types)\n","sourceCodeStart":108,"sourceCodeEnd":144,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/mesh.py#L108-L144","documentation":"Mesh/AbstractMesh __new__ normalizes the axis_types argument (defaulting to a tuple of AxisType.AUTO) and requires every element to be a jax.sharding.AxisType enum member. Passing strings like 'manual' or arbitrary objects raises TypeError via _normalize_axis_types.","triggerScenarios":"Constructing jax.sharding.Mesh(devices, axis_names, axis_types=['manual']) or passing a single non-AxisType scalar; converting code that used string sharding types from other frameworks.","commonSituations":"Newer JAX versions adding the AxisType (auto/manual/explicit) parameter; translating PyTorch/FSDP-style 'manual' string flags; LLM-training configs that specify mesh autosharding mode.","solutions":["Use jax.sharding.AxisType members: AxisType.Auto, AxisType.Manual, AxisType.Explicit","Pass a tuple of length len(axis_names), a single AxisType (broadcast), or None","Check spelling of the enum attribute against the installed JAX version's jax.sharding.AxisType"],"exampleFix":"# before\nmesh = jax.sharding.Mesh(devs, ('data', 'model'), axis_types=('auto', 'manual'))\n\n# after\nfrom jax.sharding import AxisType\nmesh = jax.sharding.Mesh(devs, ('data', 'model'), axis_types=(AxisType.Auto, AxisType.Manual))","handlingStrategy":"type-guard","validationCode":"from jax.sharding import AxisType\nok = axis_types is None or all(isinstance(a, AxisType) for a in\n    (axis_types if isinstance(axis_types, tuple) else (axis_types,)))","typeGuard":"from jax.sharding import AxisType\ndef valid_axis_types(t):\n    if t is None: return True\n    ts = t if isinstance(t, tuple) else (t,)\n    return all(isinstance(a, AxisType) for a in ts)","tryCatchPattern":"try:\n    mesh = jax.sharding.Mesh(devs, names, axis_types=axis_types)\nexcept TypeError as e:\n    if 'AxisType' in str(e):\n        mesh = jax.sharding.Mesh(devs, names)  # default Auto\n    else: raise","preventionTips":["Only use jax.sharding.AxisType enum members","Never carry string sharding modes into JAX mesh configs"],"tags":["jax","mesh","axis-types","type-validation"],"backgroundTag":"invalid-enum-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}