jax-ml/jax · error · ValueError

Unsupported mode: {mode}

Error message

Unsupported mode: {mode}

What it means

jax.nn.get_scaled_dot_general_config builds a BlockScaleConfig from a `mode` string (e.g. 'e4m3-with-e8m0-scales' style modes supported in this version). An unrecognized mode string raises this ValueError.

Source

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

        return BlockScaleConfig(
            mode='nvfp4',
            block_size=16,
            data_type=dtypes.float4_e2m1fn,
            scale_type=dtypes.float8_e4m3fn,
            global_scale=one if global_scale is None else global_scale,
            infer_only=False
        )
    elif mode == 'mxfp8':
        return BlockScaleConfig(
            mode='mxfp8',
            block_size=32,
            data_type=dtypes.float8_e4m3fn,
            scale_type=dtypes.float8_e8m0fnu,
            global_scale=None,
            infer_only=False
        )
    else:
        raise ValueError(f"Unsupported mode: {mode}")

def scaled_dot_general(
    lhs, rhs,
    dimension_numbers,
    preferred_element_type=np.float32,
    configs: list[BlockScaleConfig] | None = None,
    implementation: Literal['cudnn'] | None = None,
  ):
  r"""Scaled dot general operation.

  Performs a generalized dot product with block-scaled quantization on the
  lhs and rhs inputs. This operation extends `lax.dot_general` to support
  user-defined scaling configurations.

  Essentially, the operation follows::

      a, a_scales = quantize(lhs, configs[0])
      b, b_scales = quantize(rhs, configs[1])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect the function's docstring/source for the accepted mode literals in your JAX version and use one
2. If the format you want is unsupported, construct a BlockScaleConfig manually instead of via mode
3. Upgrade JAX if the mode exists only in a newer release

Example fix

// before
cfg = jax.nn.get_scaled_dot_general_config(mode='e5m2-e8m0')

// after
cfg = jax.nn.get_scaled_dot_general_config(mode='e4m3-e8m0')  # per docstring

# inspect supported modes from the docstring; e.g.
SUPPORTED_MODES = {'e4m3-e8m0'}  # adjust to your JAX version
if mode not in SUPPORTED_MODES: raise ValueError(f'unsupported mode {mode}')

When it happens

Trigger: Passing mode='mxfp4', mode=None, or a typo like 'e5m2-e8m0' to get_scaled_dot_general_config; using a mode string from a different JAX version.

Common situations: Experimenting with MX formats; copy-pasting config code between JAX versions where mode names changed; upgrading JAX and hitting renamed modes.

Related errors

• Unknown algorithm '{algorithm}'. Expected 'fast' or 'stable'
• Unsupported implementation option: {implementation}
• scaled_matmul requires all inputs to be 3-dimensional arrays
• scaled_matmul requires inputs a and b to have matching batch
• scaled_matmul requires scales to have matching batch (B) and

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