jax-ml/jax · error · NotImplementedError

predicate is not supported with Warpgroup lowering in jaxlib

Error message

predicate is not supported with Warpgroup lowering in jaxlib < 0.11.1

What it means

This is a temporary jaxlib version gate: passing a user predicate together with Warpgroup lowering semantics requires the arrive_dyn_expect_tx support that only exists in jaxlib 0.11.1+. On older jaxlib builds the predicated TMA copy cannot be implemented, so it raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:1143

  assert copy_params.get("swizzle") is None
  assert not copy_params.get("gmem_transform")

  if is_cp_async:
    mgpu.dialect.async_load(
        src,
        dst,
        barrier=None,  # pyrefly: ignore[bad-argument-type]
        indices=indices,
        slice_lengths=slice_lengths,
        collective=ir.ArrayAttr.get([]),
        predicate=predicate,
        oob_fill_mode=ir.IntegerAttr.get(i32, oob_mode.value),
    )
    return ()

  # TODO: Remove when the minimum jaxlib version is 0.11.1
  if has_user_predicate and not hasattr(mgpu.dialect, "arrive_dyn_expect_tx_supported"):
    raise NotImplementedError(
        "predicate is not supported with Warpgroup lowering in jaxlib < 0.11.1"
    )
  match leader_tracked:
    case CopyPartition.REPLICATED:
      leader_tracked_attr = mgpu.dialect.CopyReplicatedAttr.get()
    case CopyPartition.PARTITIONED(axis):
      leader_tracked_attr = mgpu.dialect.CopyPartitionedAttr.get(axis)
    case _:
      leader_tracked_attr = None

  assert barrier is not None
  barrier_ref = barrier.as_barrier_memref()

  if is_leader_tracked_copy:
    first_block = arith_dialect.cmpi(
        arith_dialect.CmpIPredicate.eq,
        mgpu.utils.cluster_idx(collective[0]),
        mgpu.c(0, ir.IndexType.get()),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade jaxlib (and jax) to 0.11.1 or newer.
2. If upgrading is impossible, avoid the predicate (see error 2802 alternatives: clamped indices or manual masked loads).

from jax._src.pallas.mosaic_gpu import mgpu
HAS_0111 = hasattr(mgpu.dialect, 'arrive_dyn_expect_tx_supported')
if predicate is not None and not HAS_0111:
    predicate = None  # fall back to clamped indices

When it happens

Trigger: Using a predicate/mask on a copy_gmem_to_smem under Warpgroup lowering while running jaxlib < 0.11.1 (detected by absence of mgpu.dialect.arrive_dyn_expect_tx_supported).

Common situations: Environments with pinned older jaxlib (e.g. 0.10.x) running kernels newly written to use predicated TMA copies; CI matrices with mixed jaxlib versions.

Related errors

• predicate not supported
• copy_gmem_to_smem with a barrier is only supported Hopper an
• The cp.async implementation does not support user-defined pr
• copy_gmem_to_smem without a barrier is only supported on pre
• Loading from a remote ref is only supported in jaxlib versio

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