jax-ml/jax · error · NotImplementedError

Unsupported dot algorithm: {precision}.

Error message

Unsupported dot algorithm: {precision}.

What it means

When a lax.DotAlgorithm (algorithm-based precision, not a simple PrecisionLike string) is passed to a Pallas Triton dot, the lowering only supports a small whitelist of bf16→f32 presets (and a few others handled in other cases). Any other DotAlgorithm combination raises NotImplementedError at kernel lowering.

Source

Thrown at jax/_src/pallas/triton/lowering.py:2378

    match precision:
      case lax.DotAlgorithmPreset.TF32_TF32_F32:
        input_precision = tt_dialect.InputPrecision.TF32
      case lax.DotAlgorithmPreset.TF32_TF32_F32_X3:
        input_precision = tt_dialect.InputPrecision.TF32x3
      case lax.DotAlgorithmPreset.F32_F32_F32:
        input_precision = tt_dialect.InputPrecision.IEEE
      case (
          lax.DotAlgorithmPreset.F16_F16_F16
          | lax.DotAlgorithmPreset.F16_F16_F32
          | lax.DotAlgorithmPreset.BF16_BF16_BF16
          | lax.DotAlgorithmPreset.BF16_BF16_F32
          | lax.DotAlgorithmPreset.BF16_BF16_F32_X3
          | lax.DotAlgorithmPreset.BF16_BF16_F32_X6
          | lax.DotAlgorithmPreset.BF16_BF16_F32_X9
      ):
        input_precision = None
      case _:
        raise NotImplementedError(f"Unsupported dot algorithm: {precision}.")

    assert precision.supported_lhs_types is not None
    assert precision.supported_rhs_types is not None
    a = _cast(a, a_aval.dtype, precision.supported_lhs_types[0])
    b = _cast(b, b_aval.dtype, precision.supported_rhs_types[0])
    acc_dtype = precision.accumulation_type
  elif isinstance(precision, tuple):
    a_precision, b_precision = precision
    if a_precision in _TF32_PRECISIONS or b_precision in _TF32_PRECISIONS:
      input_precision = tt_dialect.InputPrecision.TF32
    elif a_aval.dtype == jnp.float32:
      input_precision = tt_dialect.InputPrecision.IEEE
    else:
      input_precision = None

    acc_dtype = out_aval.dtype
    if acc_dtype not in (jnp.int32, jnp.float16, jnp.float64):
      acc_dtype = jnp.float32

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a plain precision string (lax.Precision.DEFAULT / HIGHEST) or None instead of a DotAlgorithm in the pallas kernel
2. If bf16 inputs, use one of the supported presets: BF16_BF16_F32, BF16_BF16_F32_X3/X6/X9
3. Raise a feature request / fall back to a non-pallas jax.lax.matmul path

Example fix

# before
acc = pl.dot(a, b, precision=lax.DotAlgorithm(lax.DotAlgorithmPreset.F16_F16_F32_X9))

# after
acc = pl.dot(a, b, precision=lax.Precision.DEFAULT)

from jax._src.lax import lax as _lax
SUPPORTED = {None, _lax.DotAlgorithmPreset.BF16_BF16_F32,
             _lax.DotAlgorithmPreset.BF16_BF16_F32_X3,
             _lax.DotAlgorithmPreset.BF16_BF16_F32_X6,
             _lax.DotAlgorithmPreset.BF16_BF16_F32_X9}
assert precision in SUPPORTED or precision is None

def is_supported_precision(p) -> bool:
    return p is None or isinstance(p, str) or str(type(p)).find('DotAlgorithm') >= 0

When it happens

Trigger: Passing an explicit lax.DotAlgorithm (e.g. lax.DotAlgorithm(f32, f32, f32, ...)) or an unsupported preset such as F16_F16_F32_X5 to a pallas dot; using precision=lax.DotAlgorithmPreset.DEFAULT on a dtype combination not covered by the supported cases.

Common situations: Porting code that tuned matmul algorithms via lax.DotAlgorithm on CUDA/XLA to a Mosaic pallas kernel; assuming XLA's algorithm knobs carry over to the Triton backend.

Related errors

• Unsupported dot precision: {precision}.
• Per-operand dot precision unsupported
• Unsupported dot precision: {precision}
• 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/9de4f0f80d38fc47. Report an issue: GitHub.