jax-ml/jax · error · NotImplementedError

unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}

Error message

unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}

What it means

The Triton lowering rule for jnp.minimum only implements floating-point and integer element types. If the inputs' dtype is neither (e.g. complex, bool, or a custom dtype), lowering fails with this NotImplementedError. It reflects Triton's arith dialect ops (minnumf/minsi/minui) which only cover those categories.

Source

Thrown at jax/_src/pallas/triton/lowering.py:1427

    lax.clamp_p: lambda min, a, max: jnp.minimum(jnp.maximum(min, a), max),
    lax.logistic_p: lambda a, accuracy: 1 / (1 + jnp.exp(-a)),
    lax.is_finite_p: lambda x: jnp.logical_and(~jnp.isnan(x), ~jnp.isinf(x)),
}

for prim, fn in _JAX_FN_MAPPING.items():
  triton_lowering_rules[prim] = lower_fun(fn, multiple_results=False)


@register_lowering(lax.min_p)
def _min_lowering_rule(ctx: LoweringRuleContext, x, y):
  # TODO(slebedev): Consider allowing customizing nan behavior.
  x_aval, y_aval = ctx.avals_in
  x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    # TODO(slebedev): Triton promotes bfloat16 to float32 and back here.
    return arith_dialect.minnumf(x, y)
  if not jnp.issubdtype(x_aval.dtype, jnp.integer):
    raise NotImplementedError(
        f"unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}"
    )
  if jnp.issubdtype(x_aval.dtype, jnp.signedinteger):
    return arith_dialect.minsi(x, y)
  else:
    return arith_dialect.minui(x, y)


@register_lowering(lax.max_p)
def _max_lowering_rule(ctx: LoweringRuleContext, x, y):
  # TODO(slebedev): Consider allowing customizing nan behavior.
  x_aval, y_aval = ctx.avals_in
  x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    # TODO(slebedev): Triton promotes bfloat16 to float32 and back here.
    return arith_dialect.maxnumf(x, y)
  if not jnp.issubdtype(x_aval.dtype, jnp.integer):
    raise NotImplementedError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Convert inputs to a supported dtype before the call: use jnp.where or replace jnp.minimum on bools with jnp.logical_and / jnp.bitwise_and
  2. For complex inputs, split into real/imag parts, apply minimum to each, and recombine, or compute the comparison manually with real()-based logic
  3. Check jax/issues for float8/complex min support in the Triton backend and upgrade JAX if support landed

Example fix

// before
m = jnp.minimum(mask_a, mask_b)  # bool inputs
// after
m = jnp.logical_and(mask_a, mask_b)
Defensive patterns

Strategy: type-guard

Validate before calling

SUPPORTED = lambda d: jnp.issubdtype(d, jnp.floating) or jnp.issubdtype(d, jnp.integer)
assert SUPPORTED(x.dtype) and SUPPORTED(y.dtype), 'minimum on Triton needs float/int'

Type guard

def triton_min_safe(x, y):
    if x.dtype == jnp.bool_ or y.dtype == jnp.bool_:
        return jnp.logical_and(x, y)
    if jnp.issubdtype(x.dtype, jnp.complexfloating):
        raise NotImplementedError('complex minimum unsupported on Triton')
    return jnp.minimum(x, y)

Try / catch

try:
    kernel = pallas.triton_compile(...)  # or jitted call
except NotImplementedError as e:
    if 'unsupported dtypes' in str(e):
        # fall back to an XLA-jitted equivalent
        out = jax.jit(jnp.minimum)(x, y)
    else:
        raise

Prevention

When it happens

Trigger: Calling jnp.minimum(x, y) inside a Pallas kernel lowered to Triton where x_aval.dtype is complex64/complex128, bool_, or a non-numeric dtype; also float8 dtypes that are not yet handled.

Common situations: Writing Pallas kernels that operate on complex numbers (e.g. FFT-like code) or boolean masks using minimum instead of logical_and; migrating code from XLA/TPU where minimum on these dtypes works; enabling float8 experiments on newer JAX versions.

Related errors

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