jax-ml/jax · error · NotImplementedError

{aval_out.dtype}

Error message

{aval_out.dtype}

What it means

The Mosaic lowering for lax.add_p dispatches on the output dtype: integers map to arith.addi, floats to arith.addf. Any other dtype category (complex, bool, extended dtypes) reaches the bare raise NotImplementedError(aval_out.dtype).

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:3473

        _dtype_to_ir_type(y_dtype))
    y = vector.broadcast(y_ty, y)
  return x, y


@register_lowering_rule(
    lax.add_p, kernel_types=[*tpu_core.CoreType], ensure_mlir_values=False
)
@register_lowering_rule(ad_util.add_any_p, ensure_mlir_values=False)
def _add_lowering_rule(ctx: LoweringRuleContext, x, y):
  x, y = _bcast(x, y, ctx.avals_in[0], ctx.avals_in[1], ctx.avals_out[0],
      ctx.lowering_context.dynamic_shape_replacement_fn,
  )
  (aval_out,) = ctx.avals_out
  if jnp.issubdtype(aval_out.dtype, jnp.integer):
    return arith.addi(x, y)
  if jnp.issubdtype(aval_out.dtype, jnp.floating):
    return arith.addf(x, y)
  raise NotImplementedError(aval_out.dtype)


class FoldingError(Exception):
  pass


def _fold(x, fuel):
  if fuel <= 0:
    raise FoldingError()
  op_name = getattr(x.owner, "name", None)
  binop_folds = {
      "arith.maxsi": max,
      "arith.minsi": min,
  }
  if op_name == "arith.constant":
    if isinstance(x.type, ir.IntegerType):
      return ir.IntegerAttr(x.owner.attributes["value"]).value
    elif isinstance(x.type, ir.FloatType):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert operands to float32/float(bfloat16) before adding: lax.convert_element_type(x, jnp.float32)
2. Replace boolean addition with logical_or
3. Do the complex arithmetic outside the kernel or via lower_fun on supported parts
4. Check Mosaic dtype support table for your TPU generation

Example fix

// before
z = x + y  # complex64 operands in kernel
// after
z = (x.real + y.real) + 1j*(x.imag + y.imag)  # or convert to float32 pairs outside kernel

import jax.numpy as jnp
def add_supported(dtype):
    return jnp.issubdtype(dtype, jnp.integer) or jnp.issubdtype(dtype, jnp.floating)

def is_kernel_safe_dtype(dt): return jnp.issubdtype(dt, jnp.integer) or jnp.issubdtype(dt, jnp.floating)

When it happens

Trigger: Adding arrays with complex dtype, bool, or a non-standard dtype inside a Pallas TPU kernel.

Common situations: Using complex64 weights inside a TPU kernel (e.g. FFT pipelines); adding boolean masks with + instead of |; custom dtypes enabled via experimental config.

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Unsigned integer dtype {aval.dtype} is not supported for dot

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