{"record":{"id":"f80769ff3ffcdfc1","repo":"jax-ml/jax","slug":"only-float32-and-int32-results-are-supported-got","errorCode":null,"errorMessage":"Only float32 and int32 results are supported, got {dtype}","messagePattern":"Only float32 and int32 results are supported, got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1435,"sourceCode":"  data in the result FIFO upon exit.\n  ```\n\n  Args:\n    shape: The shape of the result.\n    dtype: The dtype of the result.\n    mxu_index: The MXU to use.\n  \"\"\"\n  return matmul_pop_fifo_p.bind(\n      shape=shape,\n      mxu_index=mxu_index,\n      dtype=jnp.dtype(dtype),\n  )\n\n\n@matmul_pop_fifo_p.def_effectful_abstract_eval\ndef _matmul_pop_fifo_abstract_eval(*, shape, dtype, **_):\n  if dtype not in [jnp.float32, jnp.int32]:\n    raise ValueError(\n        f\"Only float32 and int32 results are supported, got {dtype}\"\n    )\n  return jax_core.ShapedArray(shape, dtype), {mxu_effect}\n\n\nconv_p = jax_core.Primitive(\"conv\")\n\n\ndef conv(\n    lhs: jax.Array,\n    rhs: jax.Array,\n    acc: jax.Array | None = None,\n    *,\n    dimension_numbers: tuple[str, str, str] | convolution.ConvDimensionNumbers,\n    window_strides: Sequence[int] | None = None,\n    padding: str | Sequence[tuple[int, int]] | None = None,\n    lhs_dilation: Sequence[int] | None = None,\n    rhs_dilation: Sequence[int] | None = None,","sourceCodeStart":1417,"sourceCodeEnd":1453,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1417-L1453","documentation":"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.","triggerScenarios":"Calling matmul_pop_fifo with dtype=jnp.bfloat16 or jnp.float16 — the result must be f32 or i32.","commonSituations":"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.","solutions":["Use jnp.float32 (or jnp.int32) for the pop_fifo result","Cast to bf16 afterwards if needed (e.g. via the output block store)"],"exampleFix":"# before\nout = matmul_pop_fifo(shape=(128,128), dtype=jnp.bfloat16)\n# after\nout = matmul_pop_fifo(shape=(128,128), dtype=jnp.float32).astype(jnp.bfloat16)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nassert dtype in (jnp.float32, jnp.int32)","typeGuard":"def supported_pop_dtype(dtype) -> bool:\n    return dtype in (jnp.float32, jnp.int32)","tryCatchPattern":null,"preventionTips":["Always pop MXU results as f32 and cast downstream"],"tags":["jax","pallas","tpu","matmul","dtype"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}