{"record":{"id":"212668c1c39e481a","repo":"jax-ml/jax","slug":"mesh-requires-the-ndim-of-its-first-argument-dev","errorCode":null,"errorMessage":"Mesh requires the ndim of its first argument (`devices`) to equal the length of its second argument (`axis_names`), but got devices.ndim == {devices.ndim} and len(axis_names) == {len(axis_names)}.","messagePattern":"Mesh requires the ndim of its first argument \\(`devices`\\) to equal the length of its second argument \\(`axis_names`\\), but got devices\\.ndim == (.+?) and len\\(axis_names\\) == (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/mesh.py","lineNumber":277,"sourceCode":"    obj = object.__new__(Mesh)\n    object.__setattr__(obj, 'devices', devices)\n    object.__setattr__(obj, 'axis_names', axis_names)\n    object.__setattr__(obj, 'axis_types', axis_types)\n    object.__setattr__(obj, 'size', size)\n    return obj\n\n  def __new__(cls, devices: np.ndarray | Sequence[xc.Device],\n              axis_names: str | Sequence[MeshAxisName],\n              axis_types: tuple[AxisType, ...] | None = None):\n    if not isinstance(devices, np.ndarray):\n      devices = np.array(devices)\n    if isinstance(axis_names, str):\n      axis_names = (axis_names,)\n    axis_names = tuple(axis_names)\n    if any(i is None for i in axis_names):\n      raise ValueError(f\"Mesh axis names cannot be None. Got: {axis_names}\")\n    if devices.ndim != len(axis_names):\n      raise ValueError(\n          \"Mesh requires the ndim of its first argument (`devices`) to equal \"\n          \"the length of its second argument (`axis_names`), but got \"\n          f\"devices.ndim == {devices.ndim} and \"\n          f\"len(axis_names) == {len(axis_names)}.\")\n\n    devices_flat = tuple(devices.flat)\n    axis_types = _normalize_axis_types(axis_names, axis_types, 'Mesh',\n                                       AxisType.Auto)\n    empty = not axis_names and devices_flat[0] is None\n    size = 0 if empty else math.prod(devices.shape)\n    return cls._create(devices_flat, devices.shape, axis_names,\n                       axis_types, size)\n\n  # No __eq__ or __hash__: interned classes use object identity.\n\n  @property\n  def is_scalar(self):\n    return self.size == 1 and not self.axis_names","sourceCodeStart":259,"sourceCodeEnd":295,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/mesh.py#L259-L295","documentation":"Mesh requires devices.ndim == len(axis_names): each mesh axis corresponds to one named dimension of the devices ndarray. If you pass a 2D device array with 3 names (or forget to reshape a flat device list), ValueError is raised.","triggerScenarios":"Mesh(jax.devices(), ('x','y')) — 1D list with two names; devices reshaped to (4,4,2) but only ('data','model') given; forgetting that a raw device list is never reshaped automatically.","commonSituations":"The single most common Mesh construction error: writing a mesh shape in the names that doesn't match the reshape applied to devices; copy-pasting mesh definitions between machines with different device counts without updating the reshape.","solutions":["Match the reshape to the names: np.array(devices).reshape(x, y) with ('x','y') names","Wrap flat devices with an explicit reshape — Mesh never reshapes for you","Compute the shape programmatically: divide jax.device_count() by known axis sizes","Use create_device_mesh to build the device array with the intended shape"],"exampleFix":"# before\nmesh = jax.sharding.Mesh(jax.devices(), ('x', 'y'))\n\n# after\nimport numpy as np\nmesh = jax.sharding.Mesh(np.array(jax.devices()).reshape(4, 8), ('x', 'y'))","handlingStrategy":"validation","validationCode":"import numpy as np\narr = np.array(devices)\nassert arr.ndim == len(axis_names), f'{arr.ndim} != {len(axis_names)}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always reshape explicitly when constructing the device array","Derive the last axis size: shape = known + (dev_count // prod(known),)","Add an assert right before Mesh construction in testable helper"],"tags":["jax","mesh","shape-mismatch","input-validation"],"backgroundTag":"ndim-length-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}