jax-ml/jax · error · ValueError

Only float32 and int32 results are supported, got {dtype}

Error message

Only float32 and int32 results are supported, got {dtype}

What it means

matmul_pop_fifo produces the matmul result as a FIFO output value, and the TPU MXU only emits float32 or int32 results. Requesting any other dtype via the shape/dtype parameters is rejected.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:1435

  data in the result FIFO upon exit.
  ```

  Args:
    shape: The shape of the result.
    dtype: The dtype of the result.
    mxu_index: The MXU to use.
  """
  return matmul_pop_fifo_p.bind(
      shape=shape,
      mxu_index=mxu_index,
      dtype=jnp.dtype(dtype),
  )


@matmul_pop_fifo_p.def_effectful_abstract_eval
def _matmul_pop_fifo_abstract_eval(*, shape, dtype, **_):
  if dtype not in [jnp.float32, jnp.int32]:
    raise ValueError(
        f"Only float32 and int32 results are supported, got {dtype}"
    )
  return jax_core.ShapedArray(shape, dtype), {mxu_effect}


conv_p = jax_core.Primitive("conv")


def conv(
    lhs: jax.Array,
    rhs: jax.Array,
    acc: jax.Array | None = None,
    *,
    dimension_numbers: tuple[str, str, str] | convolution.ConvDimensionNumbers,
    window_strides: Sequence[int] | None = None,
    padding: str | Sequence[tuple[int, int]] | None = None,
    lhs_dilation: Sequence[int] | None = None,
    rhs_dilation: Sequence[int] | None = None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jnp.float32 (or jnp.int32) for the pop_fifo result
2. Cast to bf16 afterwards if needed (e.g. via the output block store)

Example fix

# before
out = matmul_pop_fifo(shape=(128,128), dtype=jnp.bfloat16)
# after
out = matmul_pop_fifo(shape=(128,128), dtype=jnp.float32).astype(jnp.bfloat16)

import jax.numpy as jnp
assert dtype in (jnp.float32, jnp.int32)

def supported_pop_dtype(dtype) -> bool:
    return dtype in (jnp.float32, jnp.int32)

When it happens

Trigger: Calling matmul_pop_fifo with dtype=jnp.bfloat16 or jnp.float16 — the result must be f32 or i32.

Common situations: Wanting bf16 outputs to match the rest of a bf16 kernel and assuming pop_fifo supports it; copying a dtype from a VMEM block spec.

Related errors

• {ctx.avals_out[0].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/f80769ff3ffcdfc1. Report an issue: GitHub.