{"record":{"id":"fca35a36d6ec66bd","repo":"jax-ml/jax","slug":"index-map-returned-a-value-of-type-type-idx-aval","errorCode":null,"errorMessage":"index_map returned a value of type {type(idx_aval)} at position {i} with block dimension {bd} when it should be pl.Slice","messagePattern":"index_map returned a value of type (.+?) at position (.+?) with block dimension (.+?) when it should be pl\\.Slice","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":674,"sourceCode":"      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)\n              and not idx_aval.shape\n          ):\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 a scalar\"\n            )\n    for i, ov in enumerate(out_avals):\n      if ov.shape or ov.dtype not in [jnp.int32, jnp.int64]:\n        raise ValueError(\n            f\"Index map function {debug_info.func_src_info} for \"","sourceCodeStart":656,"sourceCodeEnd":692,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L656-L692","documentation":"For block dims declared as BoundedSlice, the index_map must return a pl.Slice (jax_core indexing.Slice) object at that position, because a bounded slice needs start/size info. Returning an integer or anything else for a BoundedSlice dim triggers this error.","triggerScenarios":"Using pl.BoundedSlice in block_shape but returning plain ints from index_map, or returning a Slice where an int is expected on other dims (this specific branch covers BoundedSlice dims).","commonSituations":"Adopting BoundedSlice for variable-length blocks (e.g. ragged/attention kernels) without updating the index map to produce pl.Slice instances.","solutions":["Return jax.lax.Slice or pl.slice objects for BoundedSlice dims as the docs/examples show","Return ints only for Blocked/Element/Squeezed dims","Mirror the canonical BoundedSlice example from the Pallas test suite"],"exampleFix":"# before\nspec = pl.BlockSpec(block_shape=(pl.BoundedSlice(128),),\n                    index_map=lambda i: (i,))\n# after\nspec = pl.BlockSpec(block_shape=(pl.BoundedSlice(128),),\n                    index_map=lambda i: (pl.slice(i * 128, size=128),))","handlingStrategy":"validation","validationCode":"import jax.experimental.pallas as pl\nout = index_map(*indices)\nfor v, bd in zip(out if isinstance(out, tuple) else (out,), block_shape):\n    if isinstance(bd, pl.BoundedSlice):\n        assert isinstance(v, jax.lax.Slice), 'BoundedSlice dim needs pl.Slice'","typeGuard":"def slice_for_bounded(block_shape):\n    return any(isinstance(d, pl.BoundedSlice) for d in block_shape)","tryCatchPattern":null,"preventionTips":["Follow the canonical BoundedSlice kernel examples","Keep a typed helper that builds index-map returns per dim kind"],"tags":["jax","pallas","index-map","block-spec"],"backgroundTag":"invalid-type-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}