jax-ml/jax · error · TypeError

Only support preferred_element_type in (f32, bf16, f16), but

Error message

Only support preferred_element_type in (f32, bf16, f16), but got {preferred_element_type}

What it means

The scaled_dot_general wrapper only supports float32, bfloat16, or float16 as the output (preferred) element type. Any other dtype (float64, int types, float8...) is rejected up front with a TypeError.

Source

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

  if lhs_batch == ():  # expand the last dim
    lhs_batched = jnp.expand_dims(lhs, axis=lhs.aval.ndim)
    lhs_batch = (lhs.aval.ndim,)
  if rhs_batch == ():
    rhs_batched = jnp.expand_dims(rhs, axis=rhs.aval.ndim)
    rhs_batch = (rhs.aval.ndim,)
  dn_batched = contracting_dims, (lhs_batch, rhs_batch)
  return lhs_batched, rhs_batched, dn_batched


def scaled_dot_general_wrapper(
    lhs, rhs, dimension_numbers,
    preferred_element_type=np.float32,
    configs: list[BlockScaleConfig] | None=None,
  ):
  if preferred_element_type not in (np.dtype('float32'), np.dtype('bfloat16'), np.dtype('float16')):
    msg = ('Only support preferred_element_type in (f32, bf16, f16), but got '
            '{preferred_element_type}')
    raise TypeError(msg)
  if configs is None:
    mxfp8_config = BlockScaleConfig(
        mode='mxfp8',
        block_size=32,
        data_type=dtypes.float8_e4m3fn,
        scale_type=dtypes.float8_e8m0fnu,
        global_scale=None,
        infer_only=False
    )
    configs = [mxfp8_config, mxfp8_config, mxfp8_config]

  dimension_numbers = ensure_tuple(dimension_numbers)
  lhs_batched, rhs_batched, dn_batched = _ensure_batch_dim(
      lhs, rhs, dimension_numbers
  )
  out = scaled_dot_general_fn(
      lhs_batched, rhs_batched, dn_batched, preferred_element_type, configs,
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass preferred_element_type=np.float32 (or bfloat16/float16) explicitly
2. Disable 64-bit mode if the default drifted to float64 (jax.config.update('jax_enable_x64', False))
3. Cast your inputs to f32/bf16/f16 before the call so the output dtype is naturally supported

Example fix

# before
out = wrapper(lhs, rhs, preferred_element_type=np.float64)
# after
out = wrapper(lhs, rhs, preferred_element_type=np.float32)

import numpy as np
assert preferred_element_type in (np.dtype('float32'), np.dtype('bfloat16'), np.dtype('float16'))

def is_supported_out_dtype(dt) -> bool:
    import numpy as np
    return dt in (np.dtype('float32'), np.dtype('bfloat16'), np.dtype('float16'))

When it happens

Trigger: Passing preferred_element_type=np.float64 (common when JAX is configured with x64 enabled) or an integer dtype to the cuDNN scaled matmul wrapper.

Common situations: Running with jax.config.update('jax_enable_x64', True) so defaults become float64; explicitly requesting fp64 accumulation for numerics; accidental dtype leakage from upstream arrays.

Related errors

• shardings should container 4 inputs, but got {len(shardings)
• shardings specs rank should be 3, but got lhs: {len(lhs.spec
• shardings spec for batch dim should be same, but got lhs: {l
• {ctx.avals_out[0].dtype}
• Only float32 and int32 results are supported, got {dtype}

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