{"record":{"id":"9d70b9b284ed3f74","repo":"jax-ml/jax","slug":"dimension-must-be-0-for-1d-iota","errorCode":null,"errorMessage":"Dimension must be 0 for 1D iota.","messagePattern":"Dimension must be 0 for 1D iota\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3332,"sourceCode":"  outs = []\n  for size, aval_out in zip(sizes, ctx.avals_out):\n    slice_size[axis] = size\n    outs.append(\n        vector.extract_strided_slice(\n            # pyrefly: ignore[bad-argument-type]\n            ctx.aval_to_ir_type(aval_out), x, starts, slice_size, strides\n        )\n    )\n    starts[axis] += size\n  return outs\n\n\n@register_lowering_rule(lax.iota_p)\ndef _iota_lowering_rule(ctx: LoweringRuleContext, dtype, shape, dimension,\n                        sharding):\n  if len(shape) == 1:\n    if dimension != 0:\n      raise ValueError(\"Dimension must be 0 for 1D iota.\")\n    def _1d_iota_helper():\n      iota_2d = lax.iota_p.bind(dtype=dtype,\n                                shape=(1,) + shape,\n                                dimension=1,\n                                sharding=sharding)\n      return iota_2d[0]\n    return lower_fun(_1d_iota_helper)(ctx)\n  out_type = ctx.aval_to_ir_type(ctx.avals_out[0])\n  return tpu.iota(out_type, dimensions=[dimension])\n\n\n@register_lowering_rule(lax.gather_p, kernel_types=[*tpu_core.CoreType])\ndef _gather_lowering_rule(\n    ctx: LoweringRuleContext,\n    x,\n    indices,\n    *,\n    dimension_numbers,","sourceCodeStart":3314,"sourceCodeEnd":3350,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3314-L3350","documentation":"iota on a 1-D shape can only be generated along dimension 0; requesting lax.iota with len(shape)==1 and dimension != 0 is ill-defined for the 1-D fast path and raises ValueError.","triggerScenarios":"lax.iota_p.bind (or lax.iota) with shape=(n,) and dimension=1 or any nonzero dimension in a Pallas Mosaic kernel.","commonSituations":"Kernel code constructing index vectors with iota and a dimension argument copied from a 2-D iota pattern; programmatic dimension loops hitting the degenerate 1-D case.","solutions":["Use dimension=0 for 1-D iota","For higher dimensions, build a 2-D iota and slice, as the rule itself does"],"exampleFix":"// before\nidx = lax.iota(dtype, shape=(n,), dimension=1)\n// after\nidx = lax.iota(dtype, shape=(n,), dimension=0)","handlingStrategy":"validation","validationCode":"if len(shape) == 1:\n    dimension = 0  # only valid value for 1D iota","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use dimension=0 for 1-D iota","Prefer jnp.arange for 1-D index construction in kernels"],"tags":["jax","pallas","tpu","iota","indexing"],"backgroundTag":"invalid-axis-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}