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.

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
• a and b must have the same element type, but got
• Acc ref dtype must be float32 or int32, got
• {aval_out.dtype}
• Expected an accumulator ref, got

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

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