jax-ml/jax · error · ValueError

The precision '{precision}' is not supported by dot_general

Error message

The precision '{precision}' is not supported by dot_general on CPU

What it means

Raised when lowering dot_general on the CPU platform with a precision (DotAlgorithmPreset/DotAlgorithm) outside the supported CPU set (DEFAULT, F16_F16_F16, F32_F32_F32, F64_F64_F64, BF16_BF16_F32, BF16_BF16_F32_X3, BF16_BF16_F32_X6). CPU lacks the fp8/fp4 mixed-precision units these imply.

Source

Thrown at jax/_src/lax/lax.py:6217

  # The *_ lets us reuse this for ragged_dot_general, which has group_sizes.
  lhs_aval, rhs_aval, *_ = ctx.avals_in
  lhs_dtype, rhs_dtype = lhs_aval.dtype, rhs_aval.dtype
  aval_out, = ctx.avals_out
  accumulation_aval = aval_out
  algorithm_kwarg = {}
  if isinstance(precision, (DotAlgorithm, DotAlgorithmPreset)):
    # The CPU backend silently ignores the algorithm spec, so we check here to
    # make sure that the selected algorithm is supported. We could be a little
    # bit more liberal here (any algorithm where the input and output types
    # match and all the other parameters have default values should work), but
    # it's probably sufficient to just check the presets here.
    if platform == "cpu" and precision not in {
        DotAlgorithmPreset.DEFAULT, DotAlgorithmPreset.F16_F16_F16,
        DotAlgorithmPreset.F32_F32_F32, DotAlgorithmPreset.F64_F64_F64,
        DotAlgorithmPreset.BF16_BF16_F32, DotAlgorithmPreset.BF16_BF16_F32_X3,
        DotAlgorithmPreset.BF16_BF16_F32_X6,
    }:
      raise ValueError(
          f"The precision '{precision}' is not supported by dot_general on CPU")

    # If an explicit algorithm was specified, we always cast the input types to
    # the correct types.
    def maybe_convert_dtype(operand, operand_aval, target_dtype):
      if target_dtype is None or operand_aval.dtype == target_dtype:
        return operand
      aval = core.ShapedArray(operand_aval.shape, target_dtype)
      return mlir.convert_hlo(ctx, operand, operand_aval, aval)

    lhs_dtype, rhs_dtype, accumulation_dtype = get_algorithm_compute_types(
        precision, lhs_dtype, rhs_dtype, aval_out.dtype)
    lhs = maybe_convert_dtype(lhs, lhs_aval, lhs_dtype)
    rhs = maybe_convert_dtype(rhs, rhs_aval, rhs_dtype)
    if accumulation_dtype is not None:
      accumulation_aval = core.ShapedArray(aval_out.shape, accumulation_dtype)

    if precision != DotAlgorithmPreset.DEFAULT:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Gate precision by platform: use the fp8 preset only when jax.default_backend() indicates gpu/tpu
2. Fall back to DotAlgorithmPreset.DEFAULT (or BF16 presets) on CPU
3. Ensure the process actually sees the GPU (check jax.devices()) instead of silently running on CPU

Example fix

# before
prec = DotAlgorithmPreset.F8_E4M3FN_F8_E4M3FN_F32
out = lax.dot_general(a, b, dn, precision=prec)
# after
from jax import default_backend
prec = (DotAlgorithmPreset.F8_E4M3FN_F8_E4M3FN_F32
        if default_backend() != 'cpu' else DotAlgorithmPreset.DEFAULT)
out = lax.dot_general(a, b, dn, precision=prec)

from jax import default_backend
CPU_OK = {'DEFAULT', 'F16_F16_F16', 'F32_F32_F32', 'F64_F64_F64',
          'BF16_BF16_F32', 'BF16_BF16_F32_X3', 'BF16_BF16_F32_X6'}
if default_backend() == 'cpu':
    assert precision.name in CPU_OK or precision is None, f'precision {precision} unsupported on CPU'

def precision_supported(precision, backend=None):
    from jax import default_backend
    b = backend or default_backend()
    if b != 'cpu':
        return True
    from jax._src.lax.lax import DotAlgorithmPreset as P
    return precision in {P.DEFAULT, P.F16_F16_F16, P.F32_F32_F32, P.F64_F64_F64,
                         P.BF16_BF16_F32, P.BF16_BF16_F32_X3, P.BF16_BF16_F32_X6}

try:
    out = lax.dot_general(a, b, dn, precision=prec)
except ValueError:
    out = lax.dot_general(a, b, dn, precision=DotAlgorithmPreset.DEFAULT)

When it happens

Trigger: Running lax.dot_general(..., precision=DotAlgorithmPreset.F8_E4M3FN_...) or an fp8 DotAlgorithm while the computation executes on jax.devices('cpu').

Common situations: Developing on a laptop without GPU; CI running CPU-only; device placement falling back to CPU when the GPU is busy/unavailable; forcing jax_platforms=cpu for debugging.

Related errors

• `preferred_element_type` must not be narrower than the origi
• Per-operand dot precision unsupported
• Unsupported dot precision: {precision}
• shardings should container 4 inputs, but got {len(shardings)
• shardings specs rank should be 3, but got lhs: {len(lhs.spec

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