jax-ml/jax · error · ValueError

Unrecognized mode: {config.mode}.

Error message

Unrecognized mode: {config.mode}.

What it means

The quantize helper in the cuDNN scaled matmul path only supports a fixed set of BlockScaleConfig modes (e.g. 'mxfp8', and the nvfp4-style modes). Any other config.mode string falls through to this ValueError.

Source

Thrown at jax/_src/cudnn/scaled_matmul_stablehlo.py:618

  if config.mode == "mxfp8":
    assert config.global_scale is None
    assert config.scale_type == dtypes.float8_e8m0fnu

    scales_q = cast_to_e8m0_with_rounding_up(get_scales_per_block(x))
    scaled_x = x / e8m0_to_dtype(scales_q, x.dtype)
  elif config.mode == "nvfp4":
    assert config.scale_type == dtypes.float8_e4m3fn
    assert config.global_scale.dtype == np.float32

    SCALE_MAX = dtypes.finfo(config.scale_type).max.astype(x.dtype)

    x /= config.global_scale
    scales_q = jnp.clip(get_scales_per_block(x), 0, SCALE_MAX)
    scales_q = lax.optimization_barrier(scales_q.astype(config.scale_type))
    scaled_x = x / scales_q.astype(np.float32)
  else:
    raise ValueError(f"Unrecognized mode: {config.mode}.")

  clipped_x = jnp.clip(scaled_x, -MAX, MAX)
  x_q = clipped_x.astype(config.data_type)

  x_q = x_q.reshape(x_shape)  # shape = (B, M, K)
  scales_q = jnp.reshape(scales_q, scales_q.shape[:-1]).view(
      config.scale_type
  )
  return x_q, scales_q

def scaled_dot_impl(lhs, rhs, dimension_numbers, preferred_element_type,
                    configs):
  if preferred_element_type is None:
    preferred_element_type = dtypes.result_type(
        lhs, rhs, return_weak_type_flag=False
    )
  else:
    preferred_element_type = dtypes.check_and_canonicalize_user_dtype(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a supported mode string exactly as documented (e.g. 'mxfp8', 'nf-hd', 'nvfp4-ish' per your version)
2. Upgrade jax and jax-cudnn together to the latest compatible pair
3. If you copied a config from elsewhere, compare against the modes handled in the if/elif chain above the raise

Example fix

# before
cfg = BlockScaleConfig(mode='mxf8', ...)
# after
cfg = BlockScaleConfig(mode='mxfp8', ...)

from jax._src.cudnn.scaled_matmul_stablehlo import BlockScaleConfig  # modes validated at construction in newer versions

SUPPORTED_MODES = {'mxfp8'}  # extend per your jax-cudnn version
def is_supported_mode(mode: str) -> bool:
    return mode in SUPPORTED_MODES

When it happens

Trigger: Constructing BlockScaleConfig with a mode string that is not recognized (typo like 'mxf8', or an experimental/new mode unsupported by the installed jax-cudnn version) and calling scaled_dot / scaled_dot_general which quantizes inputs.

Common situations: Upgrading jax but not jax-cudnn (or vice versa) so new modes aren't recognized; hand-rolling configs copied from other repos; typos in mode names.

Understand the failure class

Background: Invalid enum value errors: "Unknown type", "Invalid scope", "must be one of" — when a string is not on the library's allowed list — this error's family across 23 libraries.

Related errors

• invalid truth value {val!r} for environment {varname!r}
• Unrecognized config option: {name}
• For flags with a corresponding contextmanager, read their va
• Config option {name} already defined
• bool() not supported for instances of type '{0}' (did you me

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