jax-ml/jax · error · ValueError

Unexpected dim_block_size: {dim_block_size}

Error message

Unexpected dim_block_size: {dim_block_size}

What it means

When computing pointer offsets from Pallas block indices, each dimension's block specification must be either a Squeezed() dimension (lowered to constant 0) or a plain int block size (consumed from the index iterator). Any other object in the match triggers this internal ValueError, indicating an internal inconsistency or a new/unsupported BlockMapping form.

Source

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

  full_size = math.prod(full_shape) * array_dtype.itemsize
  # Use 64-bit indexing when offset might be >= 2**32 bytes.
  offset_eltype = ir.IntegerType.get_signless(64 if full_size > 2**32 else 32)
  if indexer_shape:
    offsets = _zeros(ir.RankedTensorType.get(indexer_shape, offset_eltype))
  else:
    offsets = _ir_constant(0, offset_eltype)

  indexer_iter = iter(indices)
  for dim_stride, dim_block_size, start_offset in zip(
      strides, block_info.block_shape, block_info.start_indices
  ):
    match dim_block_size:
      case pallas_core.Squeezed():
        index = _ir_constant(0, offset_eltype)
      case int():
        index = next(indexer_iter)
      case _:
        raise ValueError(f"Unexpected dim_block_size: {dim_block_size}")

    if isinstance(index, slice):
      index = primitives.Slice.from_slice(
          index, pallas_core.get_block_size(dim_block_size)
      )

    if isinstance(index, primitives.Slice):
      if index.is_dynamic_start or (index.stride != 1):
        if not index.is_dynamic_start:
          start = _ir_constant(index.start, offset_eltype)
        else:
          assert isinstance(index.start, ir.Value)
          start = index.start
        start = _ir_cast(start, offset_eltype, signed=False)

        iota = _ir_cast(
            _make_range(0, int(index.size)), offset_eltype, signed=False
        )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure jax/jaxlib versions match exactly (pip install -U jax jaxlib) and that only one pallas install is present
2. Avoid custom or experimental BlockMapping dimension specs; use standard int block sizes or squeezed dims
3. Reproduce with a minimal kernel and report the internal error to the JAX repository
4. Pin to a JAX version where your GridSpec/BlockMapping usage was known to work

try:
    compiled = kernel.compile(...)
except ValueError as e:
    if 'Unexpected dim_block_size' in str(e):
        raise RuntimeError('JAX internal version mismatch; align jax/jaxlib/pallas versions') from e
    raise

When it happens

Trigger: Passing a BlockMapping whose dimension block size is neither pallas_core.Squeezed nor an int (e.g. None, a custom sentinel, or an object from a version-mismatched pallas_core) when launching a Triton Pallas kernel; reached from load/store pointer computation.

Common situations: Mixing JAX versions (jax + jaxlib or a vendored pallas_core with mismatched types); using experimental BlockMapping features not yet supported by the Triton backend; internal API changes across JAX releases.

Related errors

• BlockMapping for {self.origin} has captured constants: {self
• Unsupported block shape type: {type(dim)}
• Unsupported device ID type: {device_id_type}
• The Pallas TPU lowering currently supports in memory space A
• Pallas TPU requires a recent libtpu version (at least 0.0.46

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