jax-ml/jax · error · NotImplementedError

Accuracy {accuracy} not supported

Error message

Accuracy {accuracy} not supported

What it means

Raised from accuracy-attribute handling for dot lowering when the given Accuracy object matches neither the exact-match form (atol/rtol/ulps producing mode='TOLERANCE' or exact) nor any other recognized branch — i.e. an unsupported accuracy specification for the target platform lowering.

Source

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

  out_type = maybe_convert_dtype(
      out_dtype, algorithm.supported_output_types(lhs_dtype, rhs_dtype)
  )
  return lhs_dtype, rhs_dtype, out_type


def accuracy_attr(accuracy) -> hlo.ResultAccuracyAttr | None:
  if accuracy is None:
    return None
  elif isinstance(accuracy, AccuracyMode):
    return hlo.ResultAccuracyAttr.get(0.0, 0.0, int(0), str(accuracy.name))
  elif isinstance(accuracy, Tolerance):
    return hlo.ResultAccuracyAttr.get(
        atol=accuracy.atol,
        rtol=accuracy.rtol,
        ulps=accuracy.ulps,
        mode='TOLERANCE',
    )
  raise NotImplementedError(f"Accuracy {accuracy} not supported")

def _handle_dot_precision(ctx, lhs, rhs, precision, platform):
  def _is_fp8_mixed_precision_matmul(_lhs_dtypes, _rhs_dtypes):
    fp8_dtypes = (dtypes.float8_e4m3fn, dtypes.float8_e5m2,
                  dtypes.float8_e5m2fnuz, dtypes.float8_e4m3fnuz,
                  dtypes.float8_e3m4, dtypes.float8_e4m3,
                  dtypes.float8_e8m0fnu)
    return _lhs_dtypes in fp8_dtypes and _rhs_dtypes in fp8_dtypes

  # 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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set atol and rtol (standard tolerance form) rather than ulps-only or exotic combos
2. Verify the Accuracy fields supported by your JAX version's lax.Accuracy docs
3. Omit accuracy if exact tolerance control is not required

# only standard tolerance combos are supported
acc = lax.Accuracy(atol=1e-3, rtol=1e-3) if accuracy is not None else None

try:
    out = lax.dot_general(a, b, dn, precision=None, accuracy=acc)
except NotImplementedError:
    out = lax.dot_general(a, b, dn)  # drop accuracy

When it happens

Trigger: Passing a lax.Accuracy instance with a combination of fields (e.g. only ulps set, or contradictory tolerances) that the HLO ResultAccuracyAttr conversion in this function cannot represent.

Common situations: Using the experimental accuracy controls on dot_general with values valid on GPU but unsupported elsewhere; version drift in the Accuracy API fields.

Related errors

• `out_sharding` passed to `dot_general` can only contain unre
• Input type is incompatible with `preferred_element_type`. Th
• `preferred_element_type` must have the same signedness as th
• `preferred_element_type` must not be narrower than the origi
• dot_general requires lhs dimension numbers to be nonnegative

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