{"record":{"id":"5d8f326e4319d155","repo":"jax-ml/jax","slug":"index-map-function-debug-info-func-src-info-for","errorCode":null,"errorMessage":"Index map function {debug_info.func_src_info} for {origin} must return {len(block_shape)} values to match {block_shape=}. Currently returning {len(unflat_avals)} values:","messagePattern":"Index map function (.+?) for (.+?) must return (.+?) values to match (.+?)\\. Currently returning (.+?) values:","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":663,"sourceCode":"      )\n\n    fake_index_map_args, fake_index_map_kwargs = \\\n        index_map_tree.unflatten([False] * index_map_tree.num_leaves)\n    debug_info = api_util.debug_info(\n        \"pallas_call index_map\",\n        index_map_func,\n        fake_index_map_args,\n        fake_index_map_kwargs,\n    )\n    with tracing_grid_env(grid, vmapped_dims):\n      closed_jaxpr, out_avals = pe.trace_to_jaxpr(\n          index_map_func,\n          ft.FTPyTree(index_map_avals, index_map_tree),\n          debug_info)\n    unflat_avals = out_avals.unflatten()\n\n    if len(unflat_avals) != len(block_shape):\n      raise ValueError(\n          f\"Index map function {debug_info.func_src_info} for \"\n          f\"{origin} must return \"\n          f\"{len(block_shape)} values to match {block_shape=}. \"\n          f\"Currently returning {len(unflat_avals)} values:\"\n      )\n    # Verify types match\n    for i, (idx_aval, bd) in enumerate(zip(unflat_avals, block_shape)):\n      match bd:\n        case BoundedSlice():\n          if not isinstance(idx_aval, indexing.Slice):\n            raise ValueError(\n                \"index_map returned a value of type\"\n                f\" {type(idx_aval)} at position {i} with block dimension\"\n                f\" {bd} when it should be pl.Slice\"\n            )\n        case Blocked() | Element() | Squeezed() | int():\n          if (\n              not isinstance(idx_aval, jax_core.ShapedArray)","sourceCodeStart":645,"sourceCodeEnd":681,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L645-L681","documentation":"The BlockSpec index_map must return exactly one index value per block dimension. to_block_mapping traces the index_map, flattens its outputs, and compares the count with len(block_shape); mismatch raises this error including the function's source info and both counts.","triggerScenarios":"index_map returning a single tuple (lambda i, j: ((i, j),)) vs multiple values, returning the wrong number of elements, or returning None/a scalar for a multi-dim block shape (e.g. block_shape=(None, 128) with index_map=lambda i: (i, 0)).","commonSituations":"Extra parentheses around tuple returns; forgetting the squeezed dims still need a placeholder index; refactoring kernels from 1D to 2D without updating the map's return arity.","solutions":["Return exactly len(block_shape) values, one per dim (use 0 for Squeezed dims)","Remove nested parentheses: return i, j not ((i, j),)","Check the counts in the error message to see which operand's map is wrong"],"exampleFix":"# before\npl.BlockSpec(block_shape=(None, 128), index_map=lambda i: (i, 0, 0))\n# after\npl.BlockSpec(block_shape=(None, 128), index_map=lambda i: (i, 0))","handlingStrategy":"validation","validationCode":"import jax\njaxpr = jax.make_jaxpr(index_map)(*grid_indices)\nn_out = len(jaxpr.out_avals)\nassert n_out == len(block_shape), f'index_map returns {n_out}, block has {len(block_shape)}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Return flat tuples: return i, j — never ((i, j),)","Count placeholders (0) for Squeezed dims when writing maps"],"tags":["jax","pallas","block-spec","index-map"],"backgroundTag":"arity-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}