{"record":{"id":"f44a92f19c159a98","repo":"jax-ml/jax","slug":"unsupported-load-reduce-operation-orig-reduce","errorCode":null,"errorMessage":"Unsupported load reduce operation: {orig_reduce}","messagePattern":"Unsupported load reduce operation: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":4166,"sourceCode":"        f\"Unimplemented transforms for TMEM refs. {transforms=}\"\n    )\n  if reduce is not None:\n    orig_reduce = reduce\n    # Make sure we will interpret the reduction as unsigned, since we erase\n    # signedness.\n    if isinstance(x_aval.dtype, jnp.unsignedinteger) and \"abs\" not in reduce:\n      reduce = \"abs\" + reduce  # type: ignore\n    match reduce:\n      case \"min\":\n        reduce_attr = mgpu.dialect.TMEMLoadReduction.Min  # pyrefly: ignore[missing-attribute]\n      case \"max\":\n        reduce_attr = mgpu.dialect.TMEMLoadReduction.Max  # pyrefly: ignore[missing-attribute]\n      case \"absmin\":\n        reduce_attr = mgpu.dialect.TMEMLoadReduction.AbsMin  # pyrefly: ignore[missing-attribute]\n      case \"absmax\":\n        reduce_attr = mgpu.dialect.TMEMLoadReduction.AbsMax  # pyrefly: ignore[missing-attribute]\n      case _:\n        raise ValueError(f\"Unsupported load reduce operation: {orig_reduce}\")\n    return tuple(mgpu.dialect.async_load_tmem(x_tmem, reduce=reduce_attr))  # type: ignore\n  return (mgpu.dialect.async_load_tmem(x_tmem),)\n\n\nwait_load_tmem_p = jax_core.Primitive(\"wait_load_tmem\")\nwait_load_tmem_p.multiple_results = True\n\ndef wait_load_tmem():\n  \"\"\"Awaits all previously asynchronous TMEM loads issued by the calling thread.\n\n  Once this function returns, the TMEM loads issued by the calling thread are\n  guaranteed to have completed. The read TMEM regions can be safely overwritten\n  by the calling thread, or any threads signalled through ``Barrier``s with\n  ``orders_tensor_core=True``.\n  \"\"\"\n  wait_load_tmem_p.bind()\n\n","sourceCodeStart":4148,"sourceCodeEnd":4184,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L4148-L4184","documentation":"The reduce argument of async_load_tmem must be one of the supported reduction names handled by the match statement (min/max variants including absmin and absmax). An unknown value raises ValueError with the original op.","triggerScenarios":"Passing reduce='mean' or a misspelled or too-new reduction name to async_load_tmem.","commonSituations":"Typos; using reduction names from a different API (e.g. jnp reduction names) or a newer JAX version's names on an older install.","solutions":["Use one of the supported names: min, max, absmin, absmax (per tcgen05 TMEMLoadReduction)","Upgrade JAX if you expect a newly added reduction op"],"exampleFix":"// before\nv, r = async_load_tmem(ref, reduce='minimum')\n// after\nv, r = async_load_tmem(ref, reduce='min')","handlingStrategy":"validation","validationCode":"assert reduce in ('min','max','absmin','absmax'), reduce","typeGuard":"def valid_reduce(op): return op in {'min','max','absmin','absmax'}","tryCatchPattern":null,"preventionTips":["Use exactly the supported reduction names; check TMEMLoadReduction for the version in use"],"tags":["jax","pallas","tmem","reduction","invalid-argument"],"backgroundTag":"invalid-enum-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}