jax-ml/jax · error · ValueError

Integer constant out of range for i64: {v}

Error message

Integer constant out of range for i64: {v}

What it means

Raised when building an i64 MLIR integer constant in the Mosaic GPU lowering layer and the Python int falls outside the int64 range. This is an internal invariant check while lowering Pallas/Mosaic GPU kernels; it indicates the JAXpr itself contained an out-of-range integer index/constant. End users essentially never trigger it through supported APIs.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:4452

  ):
    if isinstance(t, (ir.IntegerType, ir.IndexType)):
      v = int(v)
    else:
      assert isinstance(t, ir.FloatType)
      v = float(v)
    return arith_dialect.constant(t, v)
  raise NotImplementedError(f"Unsupported constant: {v!r}")


def _i32_constant(v: int) -> ir.Value:
  if v < jnp.iinfo(jnp.int32).min or v > jnp.iinfo(jnp.int32).max:
    raise ValueError(f"Integer constant out of range for i32: {v}")
  return arith_dialect.constant(ir.IntegerType.get_signless(32), v)


def _i64_constant(v: int) -> ir.Value:
  if v < jnp.iinfo(jnp.int64).min or v > jnp.iinfo(jnp.int64).max:
    raise ValueError(f"Integer constant out of range for i64: {v}")
  return arith_dialect.constant(ir.IntegerType.get_signless(64), v)


def _as_index(v: object) -> ir.Value:
  match v:
    case int():
      return arith_dialect.constant(ir.IndexType.get(), v)
    case ir.Value() if isinstance(v.type, ir.IndexType):
      return v
    case ir.Value() if isinstance(v.type, ir.IntegerType):
      return arith_dialect.index_cast(ir.IndexType.get(), v)
    case mgpu.FragmentedArray(layout=mgpu.WGSplatFragLayout()):
      return _as_index(v.registers.item())
    case jax_literals.TypedNdArray() if (
        np.issubdtype(v.dtype, np.integer) and v.ndim == 0
    ):
      return arith_dialect.constant(ir.IndexType.get(), int(v))
    case _:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect the kernel code for huge integer literals or index computations that overflow int64
2. Compute indices with numpy int64/jnp.int64 and clamp/validate before use
3. If shapes are genuinely enormous, restructure the kernel to use smaller per-block indices

Example fix

// before
offset = 2**70  # passed into kernel
// after
offset = np.int64(min(2**63 - 1, computed_offset))

import numpy as np
assert np.iinfo(np.int64).min <= v <= np.iinfo(np.int64).max, f"int64 overflow: {v}"

def is_i64(v) -> bool:
    return isinstance(v, int) and -2**63 <= v < 2**63

When it happens

Trigger: Calling a Mosaic GPU (plgpu) kernel whose computation contains an integer constant exceeding 2**63-1 or below -2**63, e.g. an enormous python int used as an index or offset inside the kernel.

Common situations: Passing Python ints produced by overflow-prone arithmetic (e.g. large shapes/strides multiplied together) into pallas kernels; bugs in user code computing indices.

Related errors

• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot specify both out_shape and out_type. Use out_type.
• Cannot specify both scratch_shapes and scratch_types. Use sc
• Unsupported dtype: {ref.dtype}

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