jax-ml/jax · error · NotImplementedError

predicate not supported

Error message

predicate not supported

What it means

copy_smem_to_gmem in the GPU interpreter does not implement the optional predicate (conditional copy) argument. Passing a non-None predicate raises NotImplementedError.

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)

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

• reduction_op not supported
• Dynamic grid bounds not (yet) supported in GPU interpret mod
• Dynamic grid bounds not (yet) supported on GPU
• Scalar arguments not (yet) supported on GPU
• Unsupported memory space: {space}

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