{"record":{"id":"0de43b3532473270","repo":"jax-ml/jax","slug":"index-map-function-debug-info-func-src-info-for-0de43b","errorCode":null,"errorMessage":"Index map function {debug_info.func_src_info} for {origin} must return integer scalars. Output[{i}] has type {ov}.","messagePattern":"Index map function (.+?) for (.+?) must return integer scalars\\. Output\\[(.+?)\\] has type (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":691,"sourceCode":"          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\n    mapping = BlockMapping(\n        block_shape=block_shape,\n        transformed_block_aval=block_aval,  # There are no transforms by default\n        index_map_jaxpr=closed_jaxpr,\n        index_map_out_tree=out_avals.tree,\n        array_aval=array_aval,\n        origin=origin,","sourceCodeStart":673,"sourceCodeEnd":709,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L673-L709","documentation":"Every value returned by a BlockSpec index_map must be an int32 or int64 scalar with empty shape. After tracing, to_block_mapping verifies each output aval's shape is empty and dtype is an integer type; otherwise it raises, naming the offending output index and its type.","triggerScenarios":"index_map returning floats (e.g. i * 0.5), booleans, or shaped arrays — e.g. lambda b: (b * step,) where step is a Python float, or returning a (1,) shaped array.","commonSituations":"Block-size arithmetic done with floats; index math producing bool from comparisons; returning indices computed in float32 for large-shape convenience.","solutions":["Coerce to integer: use int step values or i * np.int32(step), and .astype(jnp.int32) on traced values","Ensure returned avals are shape () scalars","Keep all block-offset constants as Python ints"],"exampleFix":"# before\nindex_map=lambda b: (b * (128.0 / scale),)\n# after\nstep = int(128 // scale)\nindex_map=lambda b: (b * step,)","handlingStrategy":"validation","validationCode":"import jax, numpy as np\njaxpr = jax.make_jaxpr(index_map)(*map(jax.core.dim_constant, grid))\nassert all(not oa.shape and oa.dtype in (np.int32, np.int64) for oa in jaxpr.out_avals)","typeGuard":"def integer_scalar_map(index_map, args):\n    av = jax.make_jaxpr(index_map)(*args).out_avals\n    return all(not a.shape and a.dtype in (jnp.int32, jnp.int64) for a in av)","tryCatchPattern":null,"preventionTips":["Keep tile/step constants as Python ints","Cast traced offsets with .astype(jnp.int32) before returning"],"tags":["jax","pallas","index-map","dtype"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}