jax-ml/jax · error · ValueError

Dimension must be 0 for 1D iota.

Error message

Dimension must be 0 for 1D iota.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:3332

  outs = []
  for size, aval_out in zip(sizes, ctx.avals_out):
    slice_size[axis] = size
    outs.append(
        vector.extract_strided_slice(
            # pyrefly: ignore[bad-argument-type]
            ctx.aval_to_ir_type(aval_out), x, starts, slice_size, strides
        )
    )
    starts[axis] += size
  return outs


@register_lowering_rule(lax.iota_p)
def _iota_lowering_rule(ctx: LoweringRuleContext, dtype, shape, dimension,
                        sharding):
  if len(shape) == 1:
    if dimension != 0:
      raise ValueError("Dimension must be 0 for 1D iota.")
    def _1d_iota_helper():
      iota_2d = lax.iota_p.bind(dtype=dtype,
                                shape=(1,) + shape,
                                dimension=1,
                                sharding=sharding)
      return iota_2d[0]
    return lower_fun(_1d_iota_helper)(ctx)
  out_type = ctx.aval_to_ir_type(ctx.avals_out[0])
  return tpu.iota(out_type, dimensions=[dimension])


@register_lowering_rule(lax.gather_p, kernel_types=[*tpu_core.CoreType])
def _gather_lowering_rule(
    ctx: LoweringRuleContext,
    x,
    indices,
    *,
    dimension_numbers,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use dimension=0 for 1-D iota
2. For higher dimensions, build a 2-D iota and slice, as the rule itself does

Example fix

// before
idx = lax.iota(dtype, shape=(n,), dimension=1)
// after
idx = lax.iota(dtype, shape=(n,), dimension=0)

if len(shape) == 1:
    dimension = 0  # only valid value for 1D iota

When it happens

Trigger: lax.iota_p.bind (or lax.iota) with shape=(n,) and dimension=1 or any nonzero dimension in a Pallas Mosaic kernel.

Common situations: 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.

Related errors

• Cannot do int indexing on TPU
• Can only load a single key per load.
• Integer indexing of refs that follows a non-trivial slice is
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/9d70b9b284ed3f74. Report an issue: GitHub.