jax-ml/jax · error · ValueError

`dimension` must be smaller than the rank of the array.

Error message

`dimension` must be smaller than the rank of the array.

What it means

FragmentedArray.broadcasted_iota creates an iota (index sequence) distributed along one dimension of a register-level array on GPU. The `dimension` argument selects which axis the iota counts along, so it must index into the array's shape. If dimension >= len(shape) the axis doesn't exist and the lowering cannot proceed.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:1162

    return cls(
        _registers=np.full(layout.registers_shape(shape), value, dtype=object),
        _layout=layout,
        _is_signed=is_signed,
    )

  @staticmethod
  def broadcasted_iota(
      dtype: ir.Type,
      shape: tuple[int, ...],
      dimension: int,
      layout: FragmentedLayout | None = None,
      *,
      is_signed: bool | None = None,
  ) -> FragmentedArray:
    """Creates a broadcasted iota array along the specified dimension."""
    if dimension >= len(shape):
      raise ValueError(
          "`dimension` must be smaller than the rank of the array."
      )

    def cast(idx: ir.Value) -> ir.Value:
      if isinstance(dtype, ir.FloatType):
        i32 = ir.IntegerType.get_signless(32)
        return arith.uitofp(dtype, arith.index_cast(i32, idx))
      return arith.index_cast(dtype, idx)

    return mgpu.FragmentedArray.splat(
        llvm.mlir_undef(dtype),
        shape,
        layout,
        is_signed=is_signed,
    ).foreach(
        lambda _, idx: cast(idx[dimension]),
        create_array=True,
        is_signed=is_signed,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check that 0 <= dimension < len(shape) before calling broadcasted_iota and fix the off-by-one
2. If you need a higher-rank iota, pass a shape with rank > dimension (e.g. append a size-1 axis)
3. Log shape and dimension at kernel-build time to catch rank mismatches early

Example fix

# before
fa = FragmentedArray.broadcasted_iota(i32, (16,), dimension=1)
# after
fa = FragmentedArray.broadcasted_iota(i32, (16, 1), dimension=1)

assert 0 <= dimension < len(shape), f"dimension {dimension} out of range for shape {shape}"

def valid_iota_dim(shape: tuple[int, ...], dim: int) -> bool:
    return isinstance(dim, int) and 0 <= dim < len(shape)

When it happens

Trigger: Calling FragmentedArray.broadcasted_iota(dtype, shape, dimension=...) with a dimension index greater than or equal to rank(shape), e.g. shape=(8,) with dimension=1, or passing a dimension for a scalar shape ().

Common situations: Dynamically computing the iota dimension from loop variables or user-supplied rank in a Mosaic GPU kernel; off-by-one errors when dimension is derived from len(shape); refactoring a kernel from 2-D to 1-D tiles without updating the hardcoded dimension.

Related errors

• Swizzle {self.swizzle} requires the trailing dimension to be
• grid_names must have the same length as grid, got {self}.
• cluster_names must have the same length as cluster, got {sel
• {tiling=} and {grid=} must have same length.
• The stored value has shape {src.shape}, but the target refer

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