jax-ml/jax · error · NotImplementedError

Unsupported

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.

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

• Per-operand dot precision unsupported
• Unsigned integer dtype
• Unsupported dot precision
• Acc ref dtype must be float32 or int32, got
• Acc ref must be at least 2D, got shape

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

Appendix: 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)