{"record":{"id":"804d6dcdd6092dc0","repo":"jax-ml/jax","slug":"predicate-not-supported","errorCode":null,"errorMessage":"predicate not supported","messagePattern":"predicate not supported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py","lineNumber":1410,"sourceCode":"    thread: memory.Warpgroup,\n    src_allocation_key_as_array: jax.Array,\n    src_transforms: tuple[Any, ...],\n    dst_allocation_key_as_array: jax.Array,\n    dst_transforms: tuple[Any, ...],\n    predicate: jax.Array | None,\n    source_info: source_info_util.SourceInfo,\n    commit_group: bool,\n    reduction_op: mgpu.TMAReductionOp,\n):\n  # TODO(jburnim,paulbib): Implement commit_group.\n  del commit_group\n  src_allocation_key = HostAllocationKey.from_array(src_allocation_key_as_array)\n  src_transforms = jax.tree.map(int, _remove_noop_transforms(src_transforms))\n  dst_allocation_key = HostAllocationKey.from_array(dst_allocation_key_as_array)\n  dst_transforms = jax.tree.map(int, _remove_noop_transforms(dst_transforms))\n\n  if predicate is not None:\n    raise NotImplementedError(\"predicate not supported\")\n  if reduction_op is not None:\n    raise NotImplementedError(\"reduction_op not supported\")\n\n  clock = None\n\n  shared_memory = _get_shared_memory()\n  if shared_memory.detect_races:\n    clock = shared_memory.incr_clock(thread)\n\n  task = AsyncCopySmemToGmemTask(\n      mesh_location=mesh_location,\n      thread=thread,\n      src_allocation_key=src_allocation_key,\n      src_transforms=src_transforms,\n      dst_allocation_key=dst_allocation_key,\n      dst_transforms=dst_transforms,\n      source_info=source_info,\n      clock=clock,","sourceCodeStart":1392,"sourceCodeEnd":1428,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py#L1392-L1428","documentation":"copy_smem_to_gmem in the GPU interpreter does not implement the optional predicate (conditional copy) argument. Passing a non-None predicate raises NotImplementedError.","triggerScenarios":"Calling copy_smem_to_gmem (SMEM→GMEM copy primitive) with predicate=... while running in GPU interpret mode.","commonSituations":"Using conditional DMA copies in a kernel that works on device but is then run under the interpreter for debugging or tests.","solutions":["Remove the predicate argument and guard the copy with control flow or a masked store instead","Test that code path on device rather than in interpret mode","Update jax — interpret support for predicated copies may have landed in newer versions","File a feature request if you need predicated copies in interpret mode"],"exampleFix":"# before\ncopy_smem_to_gmem(src, dst, predicate=pred)\n# after\nif flag:  # hoist condition outside interpret run or use masked stores\ncopy_smem_to_gmem(src, dst)","handlingStrategy":"fallback","validationCode":"if interpret_mode and predicate is not None:\n    raise SkipTest('predicated copy unsupported in interpret mode')","typeGuard":null,"tryCatchPattern":"try:\n    copy_smem_to_gmem(src, dst, predicate=p)\nexcept NotImplementedError:\n    fall back to non-predicated copy guarded by control flow","preventionTips":["Keep an interpret-friendly code path without predicates"],"tags":["pallas","mosaic-gpu","interpret-mode","predicate","not-implemented"],"backgroundTag":"unsupported-operation-in-interpreter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}