jax-ml/jax · error · NotImplementedError

out_dtype argument in binary_op_lowering_rule_wg

Error message

out_dtype argument in binary_op_lowering_rule_wg

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:2774

    lax.shift_left_p: partial(
        _binary_op_lowering_rule,
        impl=lambda x, y: x._pointwise(arith_dialect.shli, y),
    ),
    lax.shift_right_logical_p: partial(
        _binary_op_lowering_rule,
        impl=lambda x, y: x._pointwise(arith_dialect.shrui, y),
    ),
    lax.shift_right_arithmetic_p: partial(
        _binary_op_lowering_rule,
        impl=lambda x, y: x._pointwise(arith_dialect.shrsi, y),
    ),
  })

def _binary_op_lowering_rule_wg(
    ctx: LoweringRuleContext, x, y, *, ui_impl, si_impl, f_impl=None, **kwargs,
):
  if kwargs.get('out_dtype') is not None:
    raise NotImplementedError("out_dtype argument in binary_op_lowering_rule_wg")
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if any(aval_in.shape for aval_in in ctx.avals_in):
      raise NotImplementedError(
          "Non-scalar arithmetic is not supported in warp-level lowering.")
  x_aval, y_aval = ctx.avals_in
  [out_aval] = ctx.avals_out
  x, y = _bcast_wg(x, y, *ctx.avals_in, *ctx.avals_out)
  if jnp.issubdtype(out_aval, jnp.signedinteger):
    return si_impl(x, y)
  elif jnp.issubdtype(out_aval, jnp.integer):
    return ui_impl(x, y)
  elif f_impl is not None and jnp.issubdtype(out_aval, jnp.floating):
    return f_impl(x, y)
  else:
    raise NotImplementedError(
        f"{ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype}"
    )

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
out = x + y.astype(jnp.int64)  # triggers out_dtype kwarg in lowering
// after
out = (x.astype(jnp.int64) + y.astype(jnp.int64))  # explicit casts, no out_dtype

Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp
# inside kernel authoring: avoid ops that pass out_dtype
# assert your binary ops don't carry out_dtype kwargs before launch
def safe_binop(op, x, y):
  assert not (x.dtype != y.dtype and jnp.promote_types(x.dtype, y.dtype) not in (x.dtype, y.dtype)), 'would need out_dtype'
  return op(x, y)

Try / catch

try:
  compiled = kernel.lower(...).compile()
except NotImplementedError as e:
  if 'out_dtype' in str(e): rewrite kernel with explicit astype casts
  else: raise

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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