{"record":{"id":"04373a1051211bed","repo":"jax-ml/jax","slug":"unsupported-implementation-option-implementation","errorCode":null,"errorMessage":"Unsupported implementation option: {implementation}","messagePattern":"Unsupported implementation option: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/nn/functions.py","lineNumber":1280,"sourceCode":"          sliding_window_length=sliding_window, return_residual=return_residual,\n      )\n      if return_residual:\n        # Regardless of input layout, cudnn always returns residual with\n        # (B N T) layout.\n        out, residual = out\n        residual = jnp.transpose(residual, (0, 2, 1)).astype(out.dtype)\n        out = (out, residual)\n    case None:\n      # TODO(kaixih@nvidia) Automatically select the best backend (defaults to XLA for now).\n      out = _dot_product_attention_xla(\n          query_arr, key_arr, value_arr, bias, mask, is_causal=is_causal,\n          scale=scale_val, q_seqlen=query_seq_lengths,\n          kv_seqlen=key_value_seq_lengths,\n          local_window_size=local_window_size,\n          return_residual=return_residual,\n      )\n    case _:\n      raise ValueError(f\"Unsupported implementation option: {implementation}\")\n\n  if return_residual:\n    out, residual = out\n    return jnp.reshape(out, output_shape), jnp.reshape(residual, residual_shape)\n\n  return jnp.reshape(out, output_shape)\n\ndef scaled_matmul(\n    lhs: Array,\n    rhs: Array,\n    lhs_scales: Array,\n    rhs_scales: Array,\n    preferred_element_type: DTypeLike = np.float32,\n) -> Array:\n    r\"\"\"Scaled matrix multiplication function.\n\n    Performs block-scaled matmul of `a` and `b` using `a_scales` and `b_scales`.\n    The last dim is the contracting dim, and block size is inferred.","sourceCodeStart":1262,"sourceCodeEnd":1298,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L1262-L1298","documentation":"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.","triggerScenarios":"Passing implementation=None, '', 'triton', 'sdpa', or a typo like 'cudnn ' / 'cuDNN'; passing a value valid only in a newer/older JAX version.","commonSituations":"Copying implementation flags from PyTorch ('sdpa', 'flash') or Flax examples into JAX; version drift where the set of backends changed.","solutions":["Check the function's docstring/signature for the supported literals and use one, e.g. implementation='xla' or 'cudnn'","Omit the argument to use the default backend selection","Upgrade/downgrade JAX if you need a backend name not present in your version"],"exampleFix":"// before\njax.nn.dot_product_attention(q, k, v, implementation='sdpa')\n\n// after\njax.nn.dot_product_attention(q, k, v, implementation='cudnn')  # or omit","handlingStrategy":"validation","validationCode":"SUPPORTED = {'xla', 'cudnn'}  # check docstring for your JAX version\nimpl = impl if impl in SUPPORTED else 'xla'","typeGuard":"def is_supported_impl(s: str) -> bool: return s in {'xla', 'cudnn'}","tryCatchPattern":null,"preventionTips":["Literal-type the implementation arg in wrappers","Pin JAX version and read its docstring when using backend-selection strings"],"tags":["jax","nn","attention","enum-argument"],"backgroundTag":"invalid-enum-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}