jax-ml/jax · error · NotImplementedError

loading from a block pointer is not supported

Error message

loading from a block pointer is not supported

What it means

The Triton Pallas lowering only supports loading from scalar/tensor-of pointers. A tt pointer whose pointee is itself a RankedTensorType (a 'block pointer') has no load lowering implemented.

Source

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

    cache = tt_dialect.CacheModifier.NONE
  elif cache_modifier == ".ca" or cache_modifier == ".cg":
    cache = _STR_TO_CACHE_MODIFIER[cache_modifier]
  else:
    raise ValueError(f"unsupported cache modifier: {cache_modifier}")
  if eviction_policy is None:
    evict = tt_dialect.EvictionPolicy.NORMAL
  else:
    try:
      evict = _STR_TO_EVICTION_POLICY[eviction_policy]
    except KeyError:
      raise ValueError(
          f"unsupported eviction policy: {eviction_policy}"
      ) from None

  if _is_triton_pointer_type(ptr.type):
    ptr_type = tt_dialect.PointerType(ptr.type)
    if isinstance(ptr_type.pointee_type, ir.RankedTensorType):
      raise NotImplementedError("loading from a block pointer is not supported")

  ptr_type = _element_type(ptr.type)
  if not _is_triton_pointer_type(ptr_type):
    raise ValueError(f"unsupported pointer type: {ptr_type}")
  ptr_type = tt_dialect.PointerType(ptr_type)
  if other is not None and mask is None:
    raise ValueError("other requires mask to be provided")
  if not isinstance(ptr.type, ir.RankedTensorType):
    if other is not None and isinstance(other.type, ir.RankedTensorType):
      raise ValueError("other cannot be a block if pointer is not a block")
    if mask is not None and isinstance(mask.type, ir.RankedTensorType):
      raise ValueError("mask cannot be a block if pointer is not a block")

  pointee_type = ptr_type.pointee_type
  is_int1 = isinstance(pointee_type, ir.IntegerType) and pointee_type.width == 1
  if is_int1:
    pointee_type = ir.IntegerType.get_signless(8)
    ptr = _ir_cast(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep the pointer's pointee a scalar element type: use base_ptr + flat integer offsets rather than nested pointer tensors
2. Flatten the block so loads use a 1D pointer tensor
3. Simplify the BlockMapping/grid so pointers stay element-level; report if it still reproduces

Example fix

// before
p = some_ptr_block  # ptr<tensor<...>>
v = tt.load(p)

// after
off = compute_flat_indices(...)
v = pl.load(base_ref.at[off])  # element-level pointer

# conceptual: ensure loads use element-level pointers, not ptr<tensor<...>>
def is_element_pointer_load(ptr_ir_type) -> bool:
    pointee = tt.PointerType(ptr_ir_type).pointee_type
    return not isinstance(pointee, ir.RankedTensorType)

try:
    v = pl.load(ref, mask=m)
except NotImplementedError as e:
    if 'block pointer' in str(e):
        # flatten: load via flat indices into base reference
        v = base_ref[flat_idx]

When it happens

Trigger: Loading through a value typed as tensor<...x !tt.ptr<tensor<...>>> — i.e. a block of pointers to blocks — produced by unusual pointer arithmetic or kernel shapes inside a Triton Pallas kernel.

Common situations: Experimental kernels doing pointer-to-pointer tricks; shapes where offset broadcasting yields a pointer tensor with tensor pointee; running kernels written against newer Triton block-pointer APIs.

Related errors

• cannot cast {src} to {dst_type}
• Only 2-argument concatenate is supported.
• Only concatenate along the last dimension is supported.
• Only arguments with shape [..., 1] are supported.
• Only 2-argument stack is supported in Triton.

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