{"record":{"id":"8623bf9dcd2661d5","repo":"jax-ml/jax","slug":"predicate-is-not-supported-with-warpgroup-lowering","errorCode":null,"errorMessage":"predicate is not supported with Warpgroup lowering in jaxlib < 0.11.1","messagePattern":"predicate is not supported with Warpgroup lowering in jaxlib < 0\\.11\\.1","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1143,"sourceCode":"  assert copy_params.get(\"swizzle\") is None\n  assert not copy_params.get(\"gmem_transform\")\n\n  if is_cp_async:\n    mgpu.dialect.async_load(\n        src,\n        dst,\n        barrier=None,  # pyrefly: ignore[bad-argument-type]\n        indices=indices,\n        slice_lengths=slice_lengths,\n        collective=ir.ArrayAttr.get([]),\n        predicate=predicate,\n        oob_fill_mode=ir.IntegerAttr.get(i32, oob_mode.value),\n    )\n    return ()\n\n  # TODO: Remove when the minimum jaxlib version is 0.11.1\n  if has_user_predicate and not hasattr(mgpu.dialect, \"arrive_dyn_expect_tx_supported\"):\n    raise NotImplementedError(\n        \"predicate is not supported with Warpgroup lowering in jaxlib < 0.11.1\"\n    )\n  match leader_tracked:\n    case CopyPartition.REPLICATED:\n      leader_tracked_attr = mgpu.dialect.CopyReplicatedAttr.get()\n    case CopyPartition.PARTITIONED(axis):\n      leader_tracked_attr = mgpu.dialect.CopyPartitionedAttr.get(axis)\n    case _:\n      leader_tracked_attr = None\n\n  assert barrier is not None\n  barrier_ref = barrier.as_barrier_memref()\n\n  if is_leader_tracked_copy:\n    first_block = arith_dialect.cmpi(\n        arith_dialect.CmpIPredicate.eq,\n        mgpu.utils.cluster_idx(collective[0]),\n        mgpu.c(0, ir.IndexType.get()),","sourceCodeStart":1125,"sourceCodeEnd":1161,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1125-L1161","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Upgrade jaxlib (and jax) to 0.11.1 or newer.","If upgrading is impossible, avoid the predicate (see error 2802 alternatives: clamped indices or manual masked loads)."],"exampleFix":null,"handlingStrategy":"validation","validationCode":"from jax._src.pallas.mosaic_gpu import mgpu\nHAS_0111 = hasattr(mgpu.dialect, 'arrive_dyn_expect_tx_supported')\nif predicate is not None and not HAS_0111:\n    predicate = None  # fall back to clamped indices","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pin jaxlib >= 0.11.1 in requirements when using predicated TMA.","Feature-detect jaxlib capabilities at import time and branch kernel code accordingly."],"tags":["mosaic-gpu","pallas","jaxlib-version","predicate","tma"],"backgroundTag":"library-version-too-old","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}