{"record":{"id":"47e0719367f7a0c4","repo":"jax-ml/jax","slug":"block-shape-for-origin-block-shape-must-ha","errorCode":null,"errorMessage":"Block shape for {origin} (= {block_shape}) must have the same number of dimensions as the array shape {array_aval.shape}.","messagePattern":"Block shape for (.+?) \\(= (.+?)\\) must have the same number of dimensions as the array shape (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":609,"sourceCode":"      if not hasattr(array_aval, \"shape\"):\n        raise ValueError(\n            \"Array type must have a `shape` attribute, but got\"\n            f\" {type(array_aval)}\"\n        )\n    if self.index_map is None:\n      index_map_func = default_index_map(len(array_aval.shape))\n      index_map_dbg = api_util.debug_info(\"pallas_call index_map\",\n                                          default_index_map, (),{}\n                                          )._replace(arg_names=(\"\",) * len(index_map_avals))\n      api_util.save_wrapped_fun_debug_info(index_map_func, index_map_dbg)\n    else:\n      index_map_func = self.index_map\n    if self.block_shape is None:\n      block_shape = _canonicalize_block_shape(array_aval.shape)\n    else:\n      block_shape = _canonicalize_block_shape(self.block_shape)\n      if len(array_aval.shape) != len(block_shape):\n        raise ValueError(\n            f\"Block shape for {origin} (= {block_shape}) \"\n            \"must have the same number of dimensions as the \"\n            f\"array shape {array_aval.shape}.\"\n        )\n\n    ref_block_shape = _get_ref_block_shape(block_shape)\n    if isinstance(array_aval, jax_core.ShapedArray):\n      arr_sh = array_aval.sharding\n      ref_sharding = arr_sh.update(spec=arr_sh.spec.update(\n          partitions=tuple(arr_sh.spec)[:len(ref_block_shape)]))\n      block_array_aval = array_aval.update(\n          shape=ref_block_shape, memory_space=jax_core.MemorySpace.Device,\n          sharding=ref_sharding)\n    elif isinstance(array_aval, state_types.AbstractLinVal):\n      if not isinstance(array_aval.inner_aval, jax_core.ShapedArray):\n        raise NotImplementedError  # TODO(mattjj,sharadmv)\n      block_array_aval = array_aval.inner_aval.update(shape=ref_block_shape)\n    else:","sourceCodeStart":591,"sourceCodeEnd":627,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L591-L627","documentation":"Each dimension of a BlockSpec's block_shape maps to one dimension of the corresponding array. to_block_mapping canonicalizes the block shape and requires len(block_shape) == len(array_aval.shape); otherwise there is no way to assign blocks to array dims and it raises this error naming the origin (input name).","triggerScenarios":"BlockSpec(block_shape=(128,), index_map=...) applied to a 2D array (e.g. shape (4096, 128)), or forgetting a trailing None for a squeezed dim (block_shape=(128, 128) on 3D data).","commonSituations":"Writing 1D-style kernels then applying them to 2D batches; adding a leading batch dim to inputs without updating the BlockSpec; copy-pasting specs between operands of different rank.","solutions":["Make block_shape length equal the array ndim, using None for squeezed dims (e.g. (None, 128) for per-row blocks of a (B, 128) array)","Update index_map to return one index per block dim","Log array_aval.shape and block_shape side by side before pallas_call"],"exampleFix":"# before\nspec = pl.BlockSpec(block_shape=(128,), index_map=lambda i: (i,))\n# on array shape (4096, 128) -> error\n# after\nspec = pl.BlockSpec(block_shape=(None, 128), index_map=lambda i: (i,))","handlingStrategy":"validation","validationCode":"assert len(block_shape) == x.ndim, (\n    f'block_shape {block_shape} rank != array ndim {x.ndim}')\nspec = pl.BlockSpec(block_shape=block_shape, index_map=...)","typeGuard":"def rank_matches(x, block_shape):\n    return x.ndim == len(block_shape)","tryCatchPattern":null,"preventionTips":["Use None for squeezed dims to preserve rank","Build BlockSpecs from x.shape programmatically instead of literals"],"tags":["jax","pallas","block-spec","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}