{"record":{"id":"b45d0bba77dcb4f1","repo":"jax-ml/jax","slug":"out-dtype-argument-in-binary-op-lowering-rule-wg","errorCode":null,"errorMessage":"out_dtype argument in binary_op_lowering_rule_wg","messagePattern":"out_dtype argument in binary_op_lowering_rule_wg","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":2774,"sourceCode":"    lax.shift_left_p: partial(\n        _binary_op_lowering_rule,\n        impl=lambda x, y: x._pointwise(arith_dialect.shli, y),\n    ),\n    lax.shift_right_logical_p: partial(\n        _binary_op_lowering_rule,\n        impl=lambda x, y: x._pointwise(arith_dialect.shrui, y),\n    ),\n    lax.shift_right_arithmetic_p: partial(\n        _binary_op_lowering_rule,\n        impl=lambda x, y: x._pointwise(arith_dialect.shrsi, y),\n    ),\n  })\n\ndef _binary_op_lowering_rule_wg(\n    ctx: LoweringRuleContext, x, y, *, ui_impl, si_impl, f_impl=None, **kwargs,\n):\n  if kwargs.get('out_dtype') is not None:\n    raise NotImplementedError(\"out_dtype argument in binary_op_lowering_rule_wg\")\n  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n    if any(aval_in.shape for aval_in in ctx.avals_in):\n      raise NotImplementedError(\n          \"Non-scalar arithmetic is not supported in warp-level lowering.\")\n  x_aval, y_aval = ctx.avals_in\n  [out_aval] = ctx.avals_out\n  x, y = _bcast_wg(x, y, *ctx.avals_in, *ctx.avals_out)\n  if jnp.issubdtype(out_aval, jnp.signedinteger):\n    return si_impl(x, y)\n  elif jnp.issubdtype(out_aval, jnp.integer):\n    return ui_impl(x, y)\n  elif f_impl is not None and jnp.issubdtype(out_aval, jnp.floating):\n    return f_impl(x, y)\n  else:\n    raise NotImplementedError(\n        f\"{ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype}\"\n    )\n","sourceCodeStart":2756,"sourceCodeEnd":2792,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L2756-L2792","documentation":"Raised by the Mosaic GPU (Pallas) warpgroup lowering rule for binary arithmetic primitives when the primitive carries an out_dtype keyword. The warpgroup lowering path simply does not implement mixed-precision binary ops where the output dtype differs from the natural result, so it rejects the op at lowering time.","triggerScenarios":"Calling a lax binary op (add/sub/mul/div/pow etc.) inside a pallas_mosaic GPU kernel with an explicit out_dtype parameter (e.g. via jnp operations that specify a result dtype), so the out_dtype kwarg reaches the warpgroup lowering rule.","commonSituations":"Writing Pallas Mosaic GPU kernels that rely on dtype-promoting or dtype-forcing arithmetic (e.g. jax.lax ops with preferred_element_type such as integer_pow, xlogy-style lowering, or mixed int32/int64 accumulation) on TPU/GPU warpgroup semantics.","solutions":["Remove or avoid out_dtype/preferred_element_type in the binary op inside the kernel; cast inputs explicitly with .astype() before the op instead","Restructure the kernel to do the arithmetic in the inputs' natural dtype and convert the result afterwards","Check for a newer JAX version where warpgroup lowering supports out_dtype","Fall back to a non-Pallas JAX implementation for that operation"],"exampleFix":"// before\nout = x + y.astype(jnp.int64)  # triggers out_dtype kwarg in lowering\n// after\nout = (x.astype(jnp.int64) + y.astype(jnp.int64))  # explicit casts, no out_dtype","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\n# inside kernel authoring: avoid ops that pass out_dtype\n# assert your binary ops don't carry out_dtype kwargs before launch\ndef safe_binop(op, x, y):\n  assert not (x.dtype != y.dtype and jnp.promote_types(x.dtype, y.dtype) not in (x.dtype, y.dtype)), 'would need out_dtype'\n  return op(x, y)","typeGuard":null,"tryCatchPattern":"try:\n  compiled = kernel.lower(...).compile()\nexcept NotImplementedError as e:\n  if 'out_dtype' in str(e): rewrite kernel with explicit astype casts\n  else: raise","preventionTips":["Always cast inputs explicitly with astype instead of relying on out_dtype/preferred_element_type in Pallas kernels","Keep kernel arithmetic in a single promoted dtype"],"tags":["jax","pallas","mosaic-gpu","dtype","lowering","not-implemented"],"backgroundTag":"unsupported-dtype-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}