jax-ml/jax · error · NotImplementedError

Per-operand dot precision unsupported

Error message

Per-operand dot precision unsupported

What it means

lax dot primitives accept per-operand precision tuples like (Precision.DEFAULT, Precision.HIGH), but Mosaic TPU lowering maps dot to a single TPU contract precision attribute, so both operands must request the same precision.

Source

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

    acc = arith.constant(
        red_type, ir.DenseElementsAttr.get_splat(red_type, val)
    )
    red = vector.multi_reduction(
        ir.Attribute.parse("#vector.kind<add>"),
        arith.mulf(x, y),
        acc,
        [1]
    )
    return vector.shape_cast(out_type, red)

  tpu_dot_dims = jax_dot_dims_to_tpu_dot_dot_dims(
      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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a uniform precision: precision=lax.Precision.HIGH or a matching tuple (p, p)
2. Pass precision=None for the default

Example fix

// before
out = lax.dot_general(a, b, ..., precision=(lax.Precision.DEFAULT, lax.Precision.HIGHEST))
// after
out = lax.dot_general(a, b, ..., precision=lax.Precision.HIGHEST)

if isinstance(precision, tuple):
    assert precision[0] == precision[1], 'per-operand precision must match on Mosaic'
precision = precision[0] if isinstance(precision, tuple) else precision

When it happens

Trigger: lax.dot_general / jnp.matmul with precision=(lax.Precision.DEFAULT, lax.Precision.HIGH) (mismatched tuple) inside a Pallas Mosaic TPU kernel.

Common situations: Code tuned for TPU-via-XLA where mixed precision hints were accepted, then moved into a Pallas kernel.

Related errors

• Unsupported dot precision: {precision}
• 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/e9f80d19776ddb97. Report an issue: GitHub.