{"record":{"id":"808c9cd3e676de75","repo":"jax-ml/jax","slug":"length-of-sharding-spec-len-out-s-spec-must-b","errorCode":null,"errorMessage":"Length of sharding.spec ({len(out_s.spec)}) must be equal to aval's ndim ({ndim}). Got sharding.spec {out_s.spec}, aval.ndim {ndim} and sharding {out_s}","messagePattern":"Length of sharding\\.spec \\((.+?)\\) must be equal to aval's ndim \\((.+?)\\)\\. Got sharding\\.spec (.+?), aval\\.ndim (.+?) and sharding (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2295,"sourceCode":"\n@cache(max_size=4096,\n       trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value)\ndef get_sharding(sharding, shape):\n  \"\"\"Modifies and checks the sharding.\n\n  Some modifications/checks include:\n    * Making the length of specs the same as ndim\n    * If a mesh axis is mentioned in pspec is Auto/Manual, replace it with None\n    * Checking for len(spec)-ndim match\n    * Checking if the mesh is an AbstractMesh.\n  \"\"\"\n  ndim = len(shape)\n  if sharding is None:\n    return _empty_sharding(ndim)\n\n  out_s = _maybe_modify_sharding(sharding, ndim)\n  if len(out_s.spec) != ndim:\n    raise ValueError(\n        f\"Length of sharding.spec ({len(out_s.spec)}) must be equal to aval's\"\n        f\" ndim ({ndim}). Got sharding.spec {out_s.spec}, aval.ndim {ndim} and\"\n        f\" sharding {out_s}\")\n  if not isinstance(out_s.mesh, mesh_lib.AbstractMesh):\n    raise ValueError(\"Mesh of an aval must be an AbstractMesh. \"\n                     f\"Got {out_s.mesh} of type {type(out_s.mesh)}\")\n  _check_divisibility(out_s, shape)\n  if out_s.memory_kind is not None:\n    raise ValueError(\n        \"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:","sourceCodeStart":2277,"sourceCodeEnd":2313,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2277-L2313","documentation":"jax/_src/core.py canonicalizes a sharding for an abstract value (aval) and requires the length of sharding.spec (a PartitionSpec) to equal the array's ndim. Each element of the spec maps to one array dimension, so a spec with a different length cannot describe how the array is laid out across the mesh.","triggerScenarios":"Passing a NamedSharding whose PartitionSpec has more/fewer entries than the array's rank, e.g. NamedSharding(mesh, P('x')) applied to a 2-D array, or a spec with None placeholders omitted for trailing dims.","commonSituations":"Reusing a sharding built for one tensor on another of different rank; forgetting that P() entries must cover every dimension including None; mesh/axis refactors changing rank expectations.","solutions":["Make len(PartitionSpec) == array.ndim, adding None for unsharded dims: P('x', None) for a 2-D array","Check arr.ndim and the spec length before attaching the sharding","If reusing shardings across tensors, key them by shape/ndim"],"exampleFix":"// before\nsh = jax.sharding.NamedSharding(mesh, P('x'))  # 1 entry\nx = jax.device_put(x_2d, sh)  # x_2d.ndim == 2\n\n// after\nsh = jax.sharding.NamedSharding(mesh, P('x', None))\nx = jax.device_put(x_2d, sh)","handlingStrategy":"validation","validationCode":"def check_sharding(spec, arr):\n    assert len(spec) == arr.ndim, (len(spec), arr.ndim)","typeGuard":"def spec_matches(spec, ndim): return len(spec) == ndim","tryCatchPattern":null,"preventionTips":["Always write PartitionSpecs with one entry per dim including None","Validate spec length vs ndim in test fixtures","Key cached shardings by (shape, ndim)"],"tags":["jax","sharding","partition-spec","shape-mismatch"],"backgroundTag":"sharding-spec-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}