{"record":{"id":"93d0a3024becfe93","repo":"jax-ml/jax","slug":"cannot-pull-iota-along-dimension-dimension-with","errorCode":null,"errorMessage":"Cannot pull iota along dimension {dimension} with None block size.","messagePattern":"Cannot pull iota along dimension (.+?) with None block size\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":2303,"sourceCode":"  local_iota = jax.lax.broadcasted_iota(dtype, iota_shape, dim_)\n  return local_iota + block_idx[dimension] * _block_size(\n      block_spec.block_shape[dimension]\n  )\n\n\n@register_pull_block_spec_rule(lax.iota_p)\ndef _iota_pull_rule(\n    ctx: PullRuleContext,\n    block_transform: BlockIndexTransform,\n    *,\n    dtype: jnp.dtype,\n    dimension: int,\n    shape: tuple[int, ...],\n    sharding: jax.sharding.Sharding,\n):\n  del ctx, sharding, dtype, shape\n  if block_transform.block_shape[dimension] is None:\n    raise ValueError(\n        f'Cannot pull iota along dimension {dimension} with None block size.'\n    )\n  return []\n\n\ndef _pattern_match_lanes_to_sublanes_reshape(\n    aval_in: core.ShapedArray,\n    aval_out: core.ShapedArray,\n) -> bool:\n  # Pattern matches a reshape of the form (..., n * l) -> (..., n, l)\n  # where l is a multiple of 128.\n\n  *leading_out, last_dim_in = aval_in.shape\n  *leading_in, second_to_last_dim_out, last_dim = aval_out.shape\n  if leading_in != leading_out:\n    return False\n  if second_to_last_dim_out * last_dim != last_dim_in:\n    return False","sourceCodeStart":2285,"sourceCodeEnd":2321,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L2285-L2321","documentation":"A ValueError (not NotImplementedError) from the iota-pulling logic: when the fuser tries to materialize an iota (index/identifier array) into a kernel argument along a dimension whose block size is None (unbounded), it cannot compute per-block index offsets, so it refuses.","triggerScenarios":"Code paths that pull lax.iota / index materialization (e.g. from jnp.arange-like patterns or positional encodings) into a fused Pallas operand where block_transform.block_shape[dimension] is None along the iota's axis.","commonSituations":"Using arange/index-derived values (RoPE positions, masks) with BlockSpecs that leave the sequence axis unbounded; dynamic-shape kernels where one axis is None for flexibility; grid computed from data rather than block shapes.","solutions":["Set a concrete integer block size on the iota's dimension in the BlockSpec","Generate the iota/arange outside the kernel and pass it as a regular block-mapped input","Derive in-kernel indices from index_map arguments (program ids) instead of a materialized iota"],"exampleFix":"// before\nspec = BlockSpec((None, d), ...)  # iota pulled along axis 0 which is None\npos = jnp.arange(seq_len)  # becomes iota\n// after\nspec = BlockSpec((128, d), ...)  # concrete block size on axis 0\n","handlingStrategy":"validation","validationCode":"assert spec.block_shape[iota_axis] is not None, 'iota axis needs a concrete block size'","typeGuard":"def iota_axis_ok(block_shape, dimension) -> bool:\n    return block_shape[dimension] is not None","tryCatchPattern":"try:\n    out = fused_fn(x)\nexcept ValueError as e:\n    if 'Cannot pull iota' in str(e):\n        out = fused_fn(x, idx=jnp.arange(n))  # pass iota as explicit input\n    else:\n        raise","preventionTips":["Materialize arange/iota outside kernels with unbounded axes","Derive indices from index_map program ids","Avoid None block dims on axes used for positional data"],"tags":["jax","pallas","iota","block-spec","validation"],"backgroundTag":"unbounded-block-dimension","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}