jax-ml/jax · error · TypeError

Input type is incompatible with `preferred_element_type`. Th

Error message

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.

What it means

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.

Source

Thrown at jax/_src/lax/lax.py:5667

  aval_out, = ctx.avals_out
  out_type = mlir.aval_to_ir_type(ctx.module_context, aval_out)
  out = hlo.bitcast_convert(out_type, operand)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

mlir.register_lowering(bitcast_convert_type_p, _bitcast_convert_type_lower)


def _validate_preferred_element_type(input_dtype, preferred_element_type):
  if (dtypes.issubdtype(input_dtype, np.integer) and
      dtypes.issubdtype(preferred_element_type, np.floating)):
    # Special-case integer->float multiply. This is allowed, and also allows
    # different signedness between input and output.
    pass
  else:
    allowed_types = (np.integer, np.floating, np.complexfloating)
    if any(dtypes.issubdtype(input_dtype, t) and not
           dtypes.issubdtype(preferred_element_type, t) for t in allowed_types):
      raise TypeError("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.")
    if (dtypes.issubdtype(input_dtype, np.signedinteger) and
        not dtypes.issubdtype(preferred_element_type, np.signedinteger)):
      raise TypeError("`preferred_element_type` must have the same signedness "
                      "as the original type.")
  input_bitwidth = np.dtype(input_dtype).itemsize
  preferred_bitwidth = np.dtype(preferred_element_type).itemsize
  if preferred_bitwidth < input_bitwidth:
    raise TypeError("`preferred_element_type` must not be narrower than the "
                    "original type.")


def _dot_general_shape_rule(lhs, rhs, *, dimension_numbers, precision,
                            preferred_element_type: DTypeLike | None,
                            out_sharding):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep the dtype category the same or lower: float32 inputs -> float32/float64 preferred, not complex
2. Compute in real precision then convert the result with .astype(jnp.complex64) afterwards
3. If complex accumulation is truly needed, cast the inputs to complex before the dot
4. Check that preferred_element_type is integer/float/complex at all (not bool)

Example fix

// before
out = lax.dot_general(a_f32, b_f32, ..., preferred_element_type=jnp.complex64)

// after
out = lax.dot_general(a_f32, b_f32, ..., preferred_element_type=jnp.float32).astype(jnp.complex64)

def preferred_ok(inp, pref):
    cats = (np.integer, np.floating, np.complexfloating)
    return not any(np.issubdtype(inp, c) and not np.issubdtype(pref, c) for c in cats)

def valid_preferred(input_dtype, preferred) -> bool:
    cats = (np.integer, np.floating, np.complexfloating)
    return not any(np.issubdtype(input_dtype, c) and not np.issubdtype(preferred, c) for c in cats)

try:
    out = lax.dot_general(a, b, dn, preferred_element_type=pref)
except TypeError:
    out = lax.dot_general(a, b, dn).astype(pref)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• `preferred_element_type` must have the same signedness as th
• `preferred_element_type` must not be narrower than the origi
• `preferred_element_type` must not be narrower than the origi
• Unsupported {preferred_element_type=}
• {} does not accept dtype {}. Accepted dtypes are subtypes of

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