jax-ml/jax · error · NotImplementedError

Unsupported dot precision: {precision}

Error message

Unsupported dot precision: {precision}

What it means

Mosaic's dot lowering only understands DEFAULT, HIGH (bf16x3), and HIGHEST (fp32) contract precisions. Any other lax.Precision value (or a raw string/enum it doesn't recognize) fails here.

Source

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

      dimension_numbers, lhs_aval.shape, rhs_aval.shape
  )

  if precision is not None:
    if precision[0] != precision[1]:
      raise NotImplementedError("Per-operand dot precision unsupported")
    precision = precision[0]
  if precision is None or precision == lax.Precision.DEFAULT:
    precision_attr = None  # That's the default in Mosaic.
  elif precision == lax.Precision.HIGH:
    precision_attr = ir.Attribute.parse(
        "#tpu.contract_precision<bf16x3>"
    )
  elif precision == lax.Precision.HIGHEST:
    precision_attr = ir.Attribute.parse(
        "#tpu.contract_precision<fp32>"
    )
  else:
    raise NotImplementedError(f"Unsupported dot precision: {precision}")
  out_tile = arith.constant(
      out_type, ir.DenseElementsAttr.get_splat(out_type, val)
  )
  # Contracting second minor is to transpose the lhs. Only try fusing if it's
  # an implicit transpose.
  implicit_transpose = (ctx.avals_in[0].ndim - 2) in lhs_dims
  return tpu.matmul(
      out_type,
      x,
      y,
      out_tile,
      dimension_numbers=tpu_dot_dims,
      precision=precision_attr,
      transpose_lhs_hint=not ctx.forward_compatible
      and ctx.lowering_context.fuse_transposed_lhs_in_matmul
      and implicit_transpose,
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use lax.Precision.DEFAULT, HIGH, or HIGHEST
2. Pass precision=None

Example fix

// before
out = lax.dot(a, b, precision=some_precision)
// after
out = lax.dot(a, b, precision=lax.Precision.HIGHEST)

import jax.lax as lax
allowed = {None, lax.Precision.DEFAULT, lax.Precision.HIGH, lax.Precision.HIGHEST}
assert precision in allowed

When it happens

Trigger: Passing a lax.Precision value outside {None, DEFAULT, HIGH, HIGHEST} to dot_general in a Pallas Mosaic kernel, e.g. a custom or legacy precision enum value.

Common situations: Cross-backend kernels that enumerate all Precision members; stale precision constants after JAX version changes.

Related errors

• Per-operand dot precision unsupported
• Unsigned integer dtype {aval.dtype} is not supported for dot
• Unsupported {preferred_element_type=}
• Per-operand conv precision unsupported
• Unsupported conv precision: {precision}

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