jax-ml/jax · error · NotImplementedError

Unsupported dot precision: {precision}.

Error message

Unsupported dot precision: {precision}.

What it means

The Pallas Triton dot lowering only accepts precision arguments that are either lax.DotAlgorithm values, a supported preset, None, or one of the classic Precision enum/string values. If precision is some other object (custom class, malformed value), the lowering falls through to raise NotImplementedError('Unsupported dot precision').

Source

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

    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
  else:
    raise NotImplementedError(f"Unsupported dot precision: {precision}.")

  a_type = ir.RankedTensorType(a.type)
  b_type = ir.RankedTensorType(b.type)
  if len(a_type.shape) != 2 or len(b_type.shape) != 2:
    raise ValueError("a and b must be 2D, but got:"
                     f" {a_type.shape} and {b_type.shape}")

  m, k = a_type.shape
  _, n = b_type.shape
  if a_type.element_type == ir.F64Type.get():
    # Triton's MMAv2 fp64 path uses the m8n8k4 PTX instruction but aggregates
    # it with NumRegisters={m:2, n:1, k:4}, producing an effective m16n8k16
    # per-warp tile.  Blocks smaller than these minimums cause repM/repN/repK
    # to round to zero, corrupting the ValueTable and segfaulting the compiler.
    #   M >= 16  (2 × instrM=8)
    #   N >=  8  (1 × instrN=8)
    #   K >= 16  (4 × instrK=4)
    errors = []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass precision=None or lax.Precision.DEFAULT/HIGH/HIGHEST explicitly
2. If using DotAlgorithm, ensure it's a lax.DotAlgorithmPreset supported by this backend
3. Print type(precision) right before the kernel launch to confirm what is actually being forwarded

Example fix

# before
acc = pl.dot(a, b, precision='high')

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

import jax.lax as lax
assert precision is None or isinstance(precision, (str, lax.Precision, lax.DotAlgorithm)), type(precision)

def valid_precision(p) -> bool:
    import jax.lax as lax
    return p is None or isinstance(p, (str, lax.Precision, lax.DotAlgorithmPreset))

When it happens

Trigger: Calling pl.dot / pallas dot with an unrecognized precision object, e.g. an arbitrary string, an int, or a stale lax.Precision member removed/renamed in a JAX version change.

Common situations: Upgrading JAX where the Precision API changed and old pickled/config values no longer match; passing precision from a config file without validating against lax.Precision.

Related errors

• end must be greater than start, but got: {end} <= {start}
• unsupported cache modifier: {cache_modifier}
• unsupported eviction policy: {eviction_policy}
• unsupported pointer type: {ptr_type}
• other requires mask to be provided

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