jax-ml/jax · error · ValueError

mask cannot be a block if pointer is not a block

Error message

mask cannot be a block if pointer is not a block

What it means

For a scalar (non-block) pointer load, the mask must also be scalar. A tensor-typed (block) mask is inconsistent with a scalar load and is rejected with this ValueError in _load.

Source

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

          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(
        ptr,
        tt_dialect.PointerType.get(pointee_type, ptr_type.address_space),
        signed=False,
    )

  if other is not None:
    other = _ir_cast(other, pointee_type, signed=False)

  result = tt_dialect.load(
      ptr,
      mask=mask,
      other=other,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a scalar mask (e.g. a single boolean condition) for scalar pointer loads
2. Or perform a block load from a block reference so the mask shape matches
3. Skip the mask for unconditional scalar loads

Example fix

// before
v = pl.load(scalar_ref, mask=idx_block < n, other=0.0)

// after
v = pl.load(scalar_ref, mask=scalar_cond, other=0.0)

def check_load_shapes(ptr_is_block: bool, mask):
    if not ptr_is_block and mask is not None:
        assert getattr(mask, 'shape', ()) == (), 'mask must be scalar for scalar loads'

def mask_matches_pointer(ptr_is_block: bool, mask) -> bool:
    if mask is None:
        return True
    mask_is_block = getattr(mask, 'shape', ()) != ()
    return ptr_is_block == mask_is_block

try:
    v = pl.load(ref, mask=m, other=o)
except ValueError:
    v = pl.load(ref, mask=bool(m_scalar), other=o)

When it happens

Trigger: Passing a block-shaped boolean mask with a scalar pointer, e.g. pl.load(ref, mask=idx_block < n) where ref refers to a single element, inside a Triton Pallas kernel.

Common situations: Boundary-checking masks written for block loads reused on scalar loads; broadcasting assumptions carried over from XLA.

Related errors

• other requires mask to be provided
• other cannot be a block if pointer is not a block
• unsupported cache modifier: {cache_modifier}
• unsupported eviction policy: {eviction_policy}
• unsupported pointer type: {ptr_type}

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