{"record":{"id":"28e5269eb720bb86","repo":"jax-ml/jax","slug":"index-map-returned-a-value-of-type-type-idx-aval-28e526","errorCode":null,"errorMessage":"index_map returned a value of type {type(idx_aval)} at position {i} with block dimension {bd} when it should be a scalar","messagePattern":"index_map returned a value of type (.+?) at position (.+?) with block dimension (.+?) when it should be a scalar","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":684,"sourceCode":"          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 \"\n            f\"{origin} must return integer scalars. Output[{i}] has type \"\n            f\"{ov}.\"\n        )\n\n    if closed_jaxpr.consts and not allow_captured_consts:\n      raise ValueError(\n          f\"Index map function {debug_info.func_src_info} for \"\n          f\"{origin} must not capture constants: {closed_jaxpr.consts}\"\n      )\n","sourceCodeStart":666,"sourceCodeEnd":702,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L666-L702","documentation":"For Blocked/Element/Squeezed/int block dims, index_map must return scalar (shape-()) values. to_block_mapping checks each output aval: if it is not a ShapedArray at all (and has no shape), it raises this error at that position, since block indices must be scalars.","triggerScenarios":"index_map returning arrays/tensors (shape (1,) etc.), lists, or non-array values for scalar block dims — e.g. lambda i: (i[:, None],) or returning a Python tuple element.","commonSituations":"Vectorizing the index computation with extra brackets; returning jnp arrays of shape (1,) instead of scalars; mixing Slice and scalar positions.","solutions":["Return scalar values: use int(i) or .squeeze() / jnp scalar per dim","Match each output to its dim kind (Slice for BoundedSlice, scalar otherwise)","Print output avals if unsure by testing the map standalone under jax.make_jaxpr"],"exampleFix":"# before\nindex_map=lambda i: (i * 128 * jnp.ones(1),)\n# after\nindex_map=lambda i: (i * 128,)","handlingStrategy":"type-guard","validationCode":"out = index_map(*idx)\nflat = out if isinstance(out, tuple) else (out,)\nassert all(not getattr(v, 'shape', ()) for v in flat), 'indices must be scalars'","typeGuard":"def all_scalar_indices(vals):\n    return all(getattr(v, 'ndim', 0) == 0 for v in vals)","tryCatchPattern":null,"preventionTips":["Never return shaped arrays from index maps","Squeeze with [()] or use Python ints for static indices"],"tags":["jax","pallas","index-map","scalar"],"backgroundTag":"invalid-type-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}