{"record":{"id":"fab11e329e991070","repo":"jax-ml/jax","slug":"number-of-axis-names-should-match-the-number-of-ax","errorCode":null,"errorMessage":"Number of axis names should match the number of axis_types. Got axis_names={axis_names} and axis_types={axis_types}","messagePattern":"Number of axis names should match the number of axis_types\\. Got axis_names=(.+?) and axis_types=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/mesh.py","lineNumber":130,"sourceCode":"  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\n\nclass BaseMesh:\n  axis_names: tuple[MeshAxisName, ...]\n  shape_tuple: tuple[tuple[str, int], ...]","sourceCodeStart":112,"sourceCodeEnd":148,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/mesh.py#L112-L148","documentation":"After type-checking, _normalize_axis_types verifies that the number of axis_types equals the number of axis_names, since each mesh axis needs exactly one type. A mismatch (e.g. 3 names, 2 types) raises ValueError with both values shown.","triggerScenarios":"Mesh(devices, ('a','b','c'), axis_types=(AxisType.Auto, AxisType.Manual)) — length mismatch. Note a single non-tuple AxisType is broadcast, so this only fires for tuples of wrong length or a scalar string expanded wrongly.","commonSituations":"Adding/removing a mesh axis (e.g. inserting a 'fsdp' axis) without updating the parallel axis_types list in a training config; copy-paste between configs with different mesh shapes.","solutions":["Make len(axis_types) == len(axis_names) (or pass one AxisType scalar to broadcast)","Add AxisType.Auto for axes you don't want to specialize","Centralize names and types in one config dict so they stay in sync"],"exampleFix":"# before\nmesh = jax.sharding.Mesh(devs, ('data','fsdp','model'),\n                         axis_types=(AxisType.Auto, AxisType.Manual))\n\n# after\nmesh = jax.sharding.Mesh(devs, ('data','fsdp','model'),\n                         axis_types=(AxisType.Auto, AxisType.Auto, AxisType.Manual))","handlingStrategy":"validation","validationCode":"names = ('data','fsdp','model')\ntypes = (AxisType.Auto, AxisType.Auto, AxisType.Manual)\nassert len(types) == len(names) or not isinstance(types, tuple)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Generate types from names: [default] * len(names)","Keep names/types pairs in one config structure"],"tags":["jax","mesh","axis-types","shape-mismatch"],"backgroundTag":"length-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}