jax-ml/jax · error · NotImplementedError

{ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype

Error message

{ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype}

What it means

The binary op lowering dispatches on the output aval's dtype class: signed integer -> addi/subi/muli, unsigned integer -> same integer ops, floating -> addf/subf/mulf (only when f_impl is provided). If the output dtype matches none of these (e.g. complex, bool, bfloat16 edge cases, or an op lacking f_impl), it raises NotImplementedError naming the primitive and both input dtypes.

Source

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

    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}"
    )


for op, si_impl, ui_impl, f_impl in [
    (lax.add_p, arith_dialect.addi, arith_dialect.addi, arith_dialect.addf),
    (lax.sub_p, arith_dialect.subi, arith_dialect.subi, arith_dialect.subf),
    (lax.mul_p, arith_dialect.muli, arith_dialect.muli, arith_dialect.mulf),
    (
        lax.div_p,
        arith_dialect.divsi,
        arith_dialect.divui,
        arith_dialect.divf,
    ),
    (lax.rem_p, arith_dialect.remsi, arith_dialect.remui, arith_dialect.remf),
    (
        lax.max_p,
        arith_dialect.maxsi,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast inputs to a supported dtype (float32/float64 or int32/int64) before the op
Avoid complex dtypes inside Mosaic GPU kernels; implement the op on real/imag parts manually
Check the registration table in lowering.py to confirm the op supports your dtype class
Use a non-Pallas implementation for unsupported dtype/op combinations

Example fix

// before
z = a_complex * b_complex  # complex dtype unsupported
// after
z_re = a.real * b.real - a.imag * b.imag
z_im = a.real * b.imag + a.imag * b.real

Defensive patterns

Strategy: type-guard

Validate before calling

import jax.numpy as jnp
SUPPORTED = (jnp.integer, jnp.floating)
def ok_dtype(*xs):
  return all(jnp.issubdtype(x.dtype, SUPPORTED) for x in xs)

Type guard

def is_supported_binary_dtype(x, y):
  d = jnp.promote_types(x.dtype, y.dtype)
  return jnp.issubdtype(d, jnp.integer) or jnp.issubdtype(d, jnp.floating)

Try / catch

try:
  out = kernel(x, y)
except NotImplementedError as e:
  if 'does not support' in str(e): cast inputs to float32/int32 and retry

Prevention

Standardize kernel inputs to float32/int32
Avoid complex dtypes in Mosaic kernels

When it happens

Trigger: Applying a lax binary primitive inside a Mosaic GPU kernel to operands whose promoted dtype is complex (or another unsupported class), or a float dtype for an op registered without an f_impl, e.g. integer-only ops on floats.

Common situations: Using complex128/complex64 math or unsupported float types (f8 etc.) in Pallas kernels; applying ops like shift_left or rem to floating inputs where no f_impl exists.

Related errors

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