jax-ml/jax · error · NotImplementedError

{ctx.avals_out[0].dtype}

Error message

{ctx.avals_out[0].dtype}

What it means

While constructing the zero-valued accumulator for a dot_general, Mosaic could not build an IR constant because the output dtype is neither a supported float nor an integer MLIR type (e.g. a complex or exotic dtype). The raw dtype is surfaced in the NotImplementedError message.

Source

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

  (aval_out,) = ctx.avals_out
  out_type = ctx.aval_to_ir_type(aval_out)
  assert isinstance(out_type, ir.ShapedType)
  val_type = ir.ShapedType(out_type).element_type
  if any(
      isinstance(val_type, cls)
      for cls in [
          ir.BF16Type,
          ir.F32Type,
          ir.Float8E5M2Type,
          ir.Float8E4M3FNType,
          ir.Float8E4M3B11FNUZType,
      ]
  ):
    val = ir.FloatAttr.get(val_type, 0.0)
  elif isinstance(val_type, ir.IntegerType):
    val = ir.IntegerAttr.get(val_type, 0)
  else:
    raise NotImplementedError(ctx.avals_out[0].dtype)
  lhs_aval, rhs_aval = ctx.avals_in
  # This is really a matrix-vector product. It only looks like matrix-matrix.
  if (
      lhs_dims == (1,)
      and rhs_dims == (1,)
      and ctx.avals_in[1].shape[0] == 1
      and len(ctx.avals_in[0].shape) == 2
      and len(ctx.avals_in[1].shape) == 2
      and (
          lhs_aval.dtype != jnp.float32
          or rhs_aval.dtype != jnp.float32
      )
  ):
    if ctx.avals_in[0].shape != ctx.avals_in[1].shape:
      bcast_shape = jnp.broadcast_shapes(
          ctx.avals_in[0].shape, ctx.avals_out[0].shape
      )
      bcast_shape = ir.VectorType.get(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compute the matmul in float and convert to complex afterwards, outside or after the dot
2. Split into real/imaginary float matmuls and recombine
3. Avoid complex dtypes in Pallas Mosaic kernels entirely

Example fix

// before
z = jnp.matmul(a_c64, b_c64)
// after
re = jnp.matmul(a_c64.real.astype(jnp.float32), b_c64.real.astype(jnp.float32))
im = jnp.matmul(a_c64.imag.astype(jnp.float32), b_c64.imag.astype(jnp.float32))

import jax.numpy as jnp
assert not jnp.issubdtype(out_dtype, jnp.complexfloating), 'complex dot unsupported in Mosaic'

def is_dot_supported_dtype(dt) -> bool:
    import jax.numpy as jnp
    return jnp.issubdtype(dt, jnp.floating) or jnp.issubdtype(dt, jnp.integer)

When it happens

Trigger: dot_general (matmul) inside a Pallas Mosaic TPU kernel whose result/accumulator dtype is complex64 or another non-float/non-int type unsupported by the constant-building path.

Common situations: Complex-valued matmuls ported from CPU/GPU to Pallas TPU kernels.

Related errors

• Only float32 and int32 results are supported, got {dtype}
• Acc ref dtype must be float32 or int32, got {dtype}
• Loading from an accumulator is not supported. Use `matmul_po
• masked swap with non-32-bit data
• Reductions over unsigned integers not implemented.

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