jax-ml/jax · error · ValueError

start and end must fit in int32

Error message

start and end must fit in int32

What it means

Triton's make_range op materializes int32 ranges, so _make_range rejects start/end values >= 2**32. Any iota, arange, offset computation, or argreduce whose bounds exceed int32 range fails during lowering with this ValueError.

Source

Thrown at jax/_src/pallas/triton/lowering.py:1502

    if i != dimension:
      iota = _expand_dims(iota, i)
  return _bcast_to(iota, shape)


def _element_type(t: ir.Type) -> ir.Type:
  if isinstance(t, ir.RankedTensorType):
    return ir.RankedTensorType(t).element_type
  else:
    return t


def _make_range(start: int, end: int) -> ir.Value:
  if end <= start:
    raise ValueError(
        f"end must be greater than start, but got: {end} <= {start}"
    )
  if max(start, end) >= 2**32:
    raise ValueError("start and end must fit in int32")
  return tt_dialect.make_range(
      ir.RankedTensorType.get([end - start], ir.IntegerType.get_signless(32)),
      start,
      end,
  )


def _full(t: ir.Type, v: Any) -> ir.Value:
  element_type = _element_type(t)
  if isinstance(element_type, ir.IntegerType):
    result = arith_dialect.constant(element_type, int(v))
  elif isinstance(element_type, ir.FloatType):
    result = arith_dialect.constant(element_type, float(v))
  else:
    raise NotImplementedError

  if isinstance(t, ir.RankedTensorType):
    return tt_dialect.splat(t, result)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split the range into 32-bit-sized chunks and process in multiple blocks/steps so each make_range call stays below 2**32
2. Keep per-kernel iteration spaces under 2**32; move the large-axis reduction across multiple kernel launches or use a different algorithm (tree reduction)
3. Check that intermediate offset products (program_id * block_size) don't silently overflow before reaching _make_range

Example fix

// before
idx = jnp.arange(0, n)  # n >= 2**32
// after
# chunked over grid so each block's range fits in int32
idx = jnp.arange(block_start, block_end)  # block_end - block_start < 2**31

INT32_MAX = 2**31 - 1
assert all(0 <= start < end < 2**31 for start, end in ranges), 'range must fit in int32'

When it happens

Trigger: jnp.arange(0, N) inside a Pallas kernel with N >= 2**32; block/grid offset computations (e.g. program_id * large_stride) whose resulting start/end values overflow int32; argreduce on tensors with more than 2**32 elements along the reduced axis.

Common situations: Very large tensors (e.g. >4B element dimensions) on big-memory GPUs; integer overflow from multiplying program_id by a large block size; 64-bit index arithmetic assumed to work on TPU but not Triton.

Related errors

• end must be greater than start, but got: {end} <= {start}
• Only positional arguments are supported by debug_print on Pa
• unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}
• cannot splat a tensor
• cannot cast {src} tp {dst_type}

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