{"record":{"id":"6bde27f69e070997","repo":"jax-ml/jax","slug":"input-type-is-incompatible-with-preferred-element","errorCode":null,"errorMessage":"Input type is incompatible with `preferred_element_type`. The compatible combinations of (input_type, preferred_element_type) are (integral, integral), (integral, floating), (floating, floating), (complex, complex.","messagePattern":"Input type is incompatible with `preferred_element_type`\\. The compatible combinations of \\(input_type, preferred_element_type\\) are \\(integral, integral\\), \\(integral, floating\\), \\(floating, floating\\), \\(complex, complex\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5667,"sourceCode":"  aval_out, = ctx.avals_out\n  out_type = mlir.aval_to_ir_type(ctx.module_context, aval_out)\n  out = hlo.bitcast_convert(out_type, operand)\n  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]\n\nmlir.register_lowering(bitcast_convert_type_p, _bitcast_convert_type_lower)\n\n\ndef _validate_preferred_element_type(input_dtype, preferred_element_type):\n  if (dtypes.issubdtype(input_dtype, np.integer) and\n      dtypes.issubdtype(preferred_element_type, np.floating)):\n    # Special-case integer->float multiply. This is allowed, and also allows\n    # different signedness between input and output.\n    pass\n  else:\n    allowed_types = (np.integer, np.floating, np.complexfloating)\n    if any(dtypes.issubdtype(input_dtype, t) and not\n           dtypes.issubdtype(preferred_element_type, t) for t in allowed_types):\n      raise TypeError(\"Input type is incompatible with \"\n                      \"`preferred_element_type`. The compatible combinations \"\n                      \"of (input_type, preferred_element_type) are \"\n                      \"(integral, integral), (integral, floating), \"\n                      \"(floating, floating), (complex, complex.\")\n    if (dtypes.issubdtype(input_dtype, np.signedinteger) and\n        not dtypes.issubdtype(preferred_element_type, np.signedinteger)):\n      raise TypeError(\"`preferred_element_type` must have the same signedness \"\n                      \"as the original type.\")\n  input_bitwidth = np.dtype(input_dtype).itemsize\n  preferred_bitwidth = np.dtype(preferred_element_type).itemsize\n  if preferred_bitwidth < input_bitwidth:\n    raise TypeError(\"`preferred_element_type` must not be narrower than the \"\n                    \"original type.\")\n\n\ndef _dot_general_shape_rule(lhs, rhs, *, dimension_numbers, precision,\n                            preferred_element_type: DTypeLike | None,\n                            out_sharding):","sourceCodeStart":5649,"sourceCodeEnd":5685,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5649-L5685","documentation":"In dot_general's preferred_element_type logic, the (input, preferred) dtype-class pairs must fall in (integral,integral), (integral,floating), (floating,floating), (complex,complex) — i.e. you cannot widen to a higher category such as floating->complex or integral->complex. Violations raise TypeError listing the allowed combinations.","triggerScenarios":"lax.dot_general(..., preferred_element_type=jnp.complex64) with float32 inputs; preferred_element_type a bool or non-numeric type; floating input with integral preferred type.","commonSituations":"Trying to accumulate matmuls in complex from real inputs; passing out_dtype-style arguments to preferred_element_type; confusing preferred_element_type (same-or-wider category only) with output casting.","solutions":["Keep the dtype category the same or lower: float32 inputs -> float32/float64 preferred, not complex","Compute in real precision then convert the result with .astype(jnp.complex64) afterwards","If complex accumulation is truly needed, cast the inputs to complex before the dot","Check that preferred_element_type is integer/float/complex at all (not bool)"],"exampleFix":"// before\nout = lax.dot_general(a_f32, b_f32, ..., preferred_element_type=jnp.complex64)\n\n// after\nout = lax.dot_general(a_f32, b_f32, ..., preferred_element_type=jnp.float32).astype(jnp.complex64)","handlingStrategy":"validation","validationCode":"def preferred_ok(inp, pref):\n    cats = (np.integer, np.floating, np.complexfloating)\n    return not any(np.issubdtype(inp, c) and not np.issubdtype(pref, c) for c in cats)","typeGuard":"def valid_preferred(input_dtype, preferred) -> bool:\n    cats = (np.integer, np.floating, np.complexfloating)\n    return not any(np.issubdtype(input_dtype, c) and not np.issubdtype(preferred, c) for c in cats)","tryCatchPattern":"try:\n    out = lax.dot_general(a, b, dn, preferred_element_type=pref)\nexcept TypeError:\n    out = lax.dot_general(a, b, dn).astype(pref)","preventionTips":["Keep preferred_element_type in the same or lower dtype category","Cast results after the dot instead of requesting cross-category accumulation","Validate preferred dtype against input dtype in helper wrappers"],"tags":["jax","dot-general","preferred-element-type","dtype-validation"],"backgroundTag":"disallowed-dtype-conversion","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}