{"record":{"id":"35896041aca60394","repo":"jax-ml/jax","slug":"only-support-preferred-element-type-in-f32-bf16","errorCode":null,"errorMessage":"Only support preferred_element_type in (f32, bf16, f16), but got {preferred_element_type}","messagePattern":"Only support preferred_element_type in \\(f32, bf16, f16\\), but got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/scaled_matmul_stablehlo.py","lineNumber":809,"sourceCode":"  if lhs_batch == ():  # expand the last dim\n    lhs_batched = jnp.expand_dims(lhs, axis=lhs.aval.ndim)\n    lhs_batch = (lhs.aval.ndim,)\n  if rhs_batch == ():\n    rhs_batched = jnp.expand_dims(rhs, axis=rhs.aval.ndim)\n    rhs_batch = (rhs.aval.ndim,)\n  dn_batched = contracting_dims, (lhs_batch, rhs_batch)\n  return lhs_batched, rhs_batched, dn_batched\n\n\ndef scaled_dot_general_wrapper(\n    lhs, rhs, dimension_numbers,\n    preferred_element_type=np.float32,\n    configs: list[BlockScaleConfig] | None=None,\n  ):\n  if preferred_element_type not in (np.dtype('float32'), np.dtype('bfloat16'), np.dtype('float16')):\n    msg = ('Only support preferred_element_type in (f32, bf16, f16), but got '\n            '{preferred_element_type}')\n    raise TypeError(msg)\n  if configs is None:\n    mxfp8_config = BlockScaleConfig(\n        mode='mxfp8',\n        block_size=32,\n        data_type=dtypes.float8_e4m3fn,\n        scale_type=dtypes.float8_e8m0fnu,\n        global_scale=None,\n        infer_only=False\n    )\n    configs = [mxfp8_config, mxfp8_config, mxfp8_config]\n\n  dimension_numbers = ensure_tuple(dimension_numbers)\n  lhs_batched, rhs_batched, dn_batched = _ensure_batch_dim(\n      lhs, rhs, dimension_numbers\n  )\n  out = scaled_dot_general_fn(\n      lhs_batched, rhs_batched, dn_batched, preferred_element_type, configs,\n  )","sourceCodeStart":791,"sourceCodeEnd":827,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/scaled_matmul_stablehlo.py#L791-L827","documentation":"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.","triggerScenarios":"Passing preferred_element_type=np.float64 (common when JAX is configured with x64 enabled) or an integer dtype to the cuDNN scaled matmul wrapper.","commonSituations":"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.","solutions":["Pass preferred_element_type=np.float32 (or bfloat16/float16) explicitly","Disable 64-bit mode if the default drifted to float64 (jax.config.update('jax_enable_x64', False))","Cast your inputs to f32/bf16/f16 before the call so the output dtype is naturally supported"],"exampleFix":"# before\nout = wrapper(lhs, rhs, preferred_element_type=np.float64)\n# after\nout = wrapper(lhs, rhs, preferred_element_type=np.float32)","handlingStrategy":"validation","validationCode":"import numpy as np\nassert preferred_element_type in (np.dtype('float32'), np.dtype('bfloat16'), np.dtype('float16'))","typeGuard":"def is_supported_out_dtype(dt) -> bool:\n    import numpy as np\n    return dt in (np.dtype('float32'), np.dtype('bfloat16'), np.dtype('float16'))","tryCatchPattern":null,"preventionTips":["Avoid jax_enable_x64 in FP8 pipelines","Explicitly pass output dtype at call sites"],"tags":["jax","dtype","fp8","matmul","unsupported-type"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}