jax-ml/jax · error · NotImplementedError

Unsupported {preferred_element_type=}

Error message

Unsupported {preferred_element_type=}

What it means

When a dot_general prefers a wider accumulation dtype, Mosaic only knows how to extend float inputs to float32. If the LHS reduction dtype is not float32 (e.g. extending bf16 to something else, or an integer promotion), it raises with the offending preferred_element_type.

Source

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

      if ctx.avals_in[0].shape != bcast_shape:
        x = vector.broadcast(bcast_shape, x)
      if ctx.avals_in[1].shape != bcast_shape:
        y = vector.broadcast(bcast_shape, y)
    red_dtype = (
        preferred_element_type if preferred_element_type else lhs_aval.dtype
    )
    red_type = ctx.aval_to_ir_type(
        lhs_aval.update(shape=(lhs_aval.shape[0],), dtype=red_dtype),
    )

    if lhs_aval.dtype != red_dtype:
      lhs_type = ctx.aval_to_ir_type(
          lhs_aval.update(shape=lhs_aval.shape, dtype=red_dtype),
      )
      if red_dtype == jnp.float32:
        x = arith.extf(lhs_type, x)
      else:
        raise NotImplementedError(f"Unsupported {preferred_element_type=}")

    if rhs_aval.dtype != red_dtype:
      rhs_type = ctx.aval_to_ir_type(
          rhs_aval.update(shape=rhs_aval.shape, dtype=red_dtype),
      )
      if red_dtype == jnp.float32:
        y = arith.extf(rhs_type, y)
      else:
        raise NotImplementedError(f"Unsupported {preferred_element_type=}")

    acc = arith.constant(
        red_type, ir.DenseElementsAttr.get_splat(red_type, val)
    )
    red = vector.multi_reduction(
        ir.Attribute.parse("#vector.kind<add>"),
        arith.mulf(x, y),
        acc,
        [1]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop preferred_element_type and let Mosaic use default accumulation
2. Use preferred_element_type=jnp.float32 with bf16/f16 inputs
3. Do the widening manually: compute at input dtype then astype the result

Example fix

// before
out = lax.dot_general(a_bf16, b_bf16, ..., preferred_element_type=jnp.int32)
// after
out = lax.dot_general(a_bf16, b_bf16, ..., preferred_element_type=jnp.float32)

import jax.numpy as jnp
preferred = None if preferred is None or preferred == jnp.float32 else jnp.float32  # normalize before dot

When it happens

Trigger: Calling lax.dot_general with preferred_element_type set to something other than float32 while the lhs is lower-precision (e.g. bf16 lhs with a non-f32 preferred type), inside a Pallas Mosaic kernel.

Common situations: Setting preferred_element_type=int32 or float64 for mixed-precision matmuls on TPU; mirroring XLA-only accumulation tricks in Pallas.

Related errors

• Unsigned integer dtype {aval.dtype} is not supported for dot
• Per-operand dot precision unsupported
• Unsupported dot precision: {precision}
• Input type is incompatible with `preferred_element_type`. Th
• {name} argument type error: {lhs.dtype}, {rhs.dtype}

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