jax-ml/jax · error · NotImplementedError

predicate not supported

Error message

predicate not supported

What it means

Capability guard in the interpreted GPU wgmma/TMA callback: a non-None predicate was supplied, but predicated TMA operations are not implemented in the interpreter, so it raises NotImplementedError. The input at fault is the predicate argument of the TMA store.

Solutions

1. Remove the predicate argument and guard the copy with control flow or a masked store instead
2. Test that code path on device rather than in interpret mode
3. Update jax — interpret support for predicated copies may have landed in newer versions
4. File a feature request if you need predicated copies in interpret mode

Example fix

# before
copy_smem_to_gmem(src, dst, predicate=pred)
# after
if flag:  # hoist condition outside interpret run or use masked stores
copy_smem_to_gmem(src, dst)

if interpret_mode and predicate is not None:
    raise SkipTest('predicated copy unsupported in interpret mode')

try:
    copy_smem_to_gmem(src, dst, predicate=p)
except NotImplementedError:
    fall back to non-predicated copy guarded by control flow

When it happens

Trigger: Calling copy_smem_to_gmem (SMEM→GMEM copy primitive) with predicate=... while running in GPU interpret mode.

Common situations: Using conditional DMA copies in a kernel that works on device but is then run under the interpreter for debugging or tests.

Related errors

• Dynamic grid bounds not (yet) supported in GPU interpret…
• Dynamic grid bounds not (yet) supported on GPU
• reduction_op not supported
• Scalar arguments not (yet) supported on GPU
• The cp.async implementation does not support user-defined…

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

Appendix: source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py:1410

    thread: memory.Warpgroup,
    src_allocation_key_as_array: jax.Array,
    src_transforms: tuple[Any, ...],
    dst_allocation_key_as_array: jax.Array,
    dst_transforms: tuple[Any, ...],
    predicate: jax.Array | None,
    source_info: source_info_util.SourceInfo,
    commit_group: bool,
    reduction_op: mgpu.TMAReductionOp,
):
  # TODO(jburnim,paulbib): Implement commit_group.
  del commit_group
  src_allocation_key = HostAllocationKey.from_array(src_allocation_key_as_array)
  src_transforms = jax.tree.map(int, _remove_noop_transforms(src_transforms))
  dst_allocation_key = HostAllocationKey.from_array(dst_allocation_key_as_array)
  dst_transforms = jax.tree.map(int, _remove_noop_transforms(dst_transforms))

  if predicate is not None:
    raise NotImplementedError("predicate not supported")
  if reduction_op is not None:
    raise NotImplementedError("reduction_op not supported")

  clock = None

  shared_memory = _get_shared_memory()
  if shared_memory.detect_races:
    clock = shared_memory.incr_clock(thread)

  task = AsyncCopySmemToGmemTask(
      mesh_location=mesh_location,
      thread=thread,
      src_allocation_key=src_allocation_key,
      src_transforms=src_transforms,
      dst_allocation_key=dst_allocation_key,
      dst_transforms=dst_transforms,
      source_info=source_info,
      clock=clock,

View on GitHub (pinned to 1e1c6a8fc0)