jax-ml/jax · error · ValueError

Unsupported implementation option: {implementation}

Error message

Unsupported implementation option: {implementation}

What it means

jax.nn.dot_product_attention dispatches on the `implementation` string ('xla', 'cudnn', 'flash' as supported in this version). Any other value falls through the match and raises this ValueError.

Source

Thrown at jax/_src/nn/functions.py:1280

          sliding_window_length=sliding_window, return_residual=return_residual,
      )
      if return_residual:
        # Regardless of input layout, cudnn always returns residual with
        # (B N T) layout.
        out, residual = out
        residual = jnp.transpose(residual, (0, 2, 1)).astype(out.dtype)
        out = (out, residual)
    case None:
      # TODO(kaixih@nvidia) Automatically select the best backend (defaults to XLA for now).
      out = _dot_product_attention_xla(
          query_arr, key_arr, value_arr, bias, mask, is_causal=is_causal,
          scale=scale_val, q_seqlen=query_seq_lengths,
          kv_seqlen=key_value_seq_lengths,
          local_window_size=local_window_size,
          return_residual=return_residual,
      )
    case _:
      raise ValueError(f"Unsupported implementation option: {implementation}")

  if return_residual:
    out, residual = out
    return jnp.reshape(out, output_shape), jnp.reshape(residual, residual_shape)

  return jnp.reshape(out, output_shape)

def scaled_matmul(
    lhs: Array,
    rhs: Array,
    lhs_scales: Array,
    rhs_scales: Array,
    preferred_element_type: DTypeLike = np.float32,
) -> Array:
    r"""Scaled matrix multiplication function.

    Performs block-scaled matmul of `a` and `b` using `a_scales` and `b_scales`.
    The last dim is the contracting dim, and block size is inferred.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check the function's docstring/signature for the supported literals and use one, e.g. implementation='xla' or 'cudnn'
2. Omit the argument to use the default backend selection
3. Upgrade/downgrade JAX if you need a backend name not present in your version

Example fix

// before
jax.nn.dot_product_attention(q, k, v, implementation='sdpa')

// after
jax.nn.dot_product_attention(q, k, v, implementation='cudnn')  # or omit

SUPPORTED = {'xla', 'cudnn'}  # check docstring for your JAX version
impl = impl if impl in SUPPORTED else 'xla'

def is_supported_impl(s: str) -> bool: return s in {'xla', 'cudnn'}

When it happens

Trigger: Passing implementation=None, '', 'triton', 'sdpa', or a typo like 'cudnn ' / 'cuDNN'; passing a value valid only in a newer/older JAX version.

Common situations: Copying implementation flags from PyTorch ('sdpa', 'flash') or Flax examples into JAX; version drift where the set of backends changed.

Related errors

• Unknown algorithm '{algorithm}'. Expected 'fast' or 'stable'
• {name} ndim should be {len(shape)}, but got {t.ndim}
• {name} dtype should be {dtype}, but got {t.dtype}
• {name} shape should be {shape}: but got {t.shape}
• The number of query heads must be a multiple of key/value he

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