jax-ml/jax · error · NotImplementedError

Unsupported dtype {x_aval.dtype}

Error message

Unsupported dtype {x_aval.dtype}

What it means

In integer_pow's repeated-squaring loop, the multiplication op is chosen from the input dtype: integer -> arith.muli, floating -> arith.mulf. Any other dtype class (complex, bool) has no multiplier and raises 'Unsupported dtype {x_aval.dtype}'.

Source

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

@register_lowering_rule(lax.integer_pow_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.integer_pow_p, mgpu.LoweringSemantics.Warpgroup)
def _integer_pow_lowering_rule(ctx: LoweringRuleContext, x, y):
  [x_aval] = ctx.avals_in
  if y == -1:
    return _lower_fun(lambda x: 1 / x)(ctx, x)
  if y <= 1:
    raise NotImplementedError

  mul_op: Callable[[Any, Any], Any]
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    mul_op = operator.mul
  elif jnp.issubdtype(x_aval.dtype, jnp.integer):
    mul_op = arith_dialect.muli
  elif jnp.issubdtype(x_aval.dtype, jnp.floating):
    mul_op = arith_dialect.mulf
  else:
    raise NotImplementedError(f"Unsupported dtype {x_aval.dtype}")

  # Y is an integer. Here we start with res = x so the range is y-1
  res = x
  # Repeated doubling algorithm.
  for i in reversed(range(y.bit_length() - 1)):
    res = mul_op(res, res)  # pyrefly: ignore[no-matching-overload]
    if (y >> i) & 1:
      res = mul_op(res, x)
  return res


@register_lowering_rule(lax.clamp_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.clamp_p, mgpu.LoweringSemantics.Warpgroup)
def _clamp_lowering_rule(ctx: LoweringRuleContext, l, x, u):
  return _lower_fun(lambda l, x, u: lax.min(lax.max(x, l), u))(ctx, l, x, u)


@register_lowering_rule(lax.square_p, mgpu.LoweringSemantics.Lane)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Expand the power manually into real/imaginary multiply-adds
2. Cast/keep inputs float or integer for power operations
3. Compute complex powers outside the kernel with standard JAX

Example fix

// before
z2 = z ** 2  # complex -> Unsupported dtype
// after
z2_re = z.real*z.real - z.imag*z.imag
z2_im = 2*z.real*z.imag

import jax.numpy as jnp
# only int/float bases support repeated-squaring powers
def pow_dtype_ok(x):
  return jnp.issubdtype(x.dtype, jnp.integer) or jnp.issubdtype(x.dtype, jnp.floating)

def supports_integer_pow(x):
  return (jnp.issubdtype(x.dtype, jnp.integer)
          or jnp.issubdtype(x.dtype, jnp.floating))

try:
  out = kernel(x)
except NotImplementedError as e:
  if 'Unsupported dtype' in str(e): expand complex power manually

When it happens

Trigger: Calling lax.integer_pow with y >= 2 on complex (or other non-int/float) operands inside a Mosaic GPU kernel.

Common situations: Computing integer powers of complex numbers (e.g. z**2, z**3) in a Pallas kernel; complex arithmetic lacks a direct MLIR mul mapping here.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Sparse metadata format not implemented for {operand_dtype=}

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