jax-ml/jax · error · TypeError

`preferred_element_type` must have the same signedness as th

Error message

`preferred_element_type` must have the same signedness as the original type.

What it means

For integral inputs, dot_general's preferred_element_type must preserve signedness: a signed input cannot prefer an unsigned accumulation type and vice versa. Mixing signed and unsigned in the accumulation would silently change semantics, so TypeError is raised.

Source

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

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):
  if out_sharding is not None and not isinstance(out_sharding, NamedSharding):
    raise NotImplementedError
  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = _from_maybe_ragged(dimension_numbers)
  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, lhs.ndim))
             for d in (lhs_contracting, lhs_batch)):
    msg = ("dot_general requires lhs dimension numbers to be nonnegative and "
           "less than the number of axes of the lhs value, got "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match signedness: signed input -> signed preferred type, unsigned -> unsigned
2. If extra range is needed, widen within the same signedness (int32 -> int64)
3. Cast the inputs to the desired signedness before the dot if overflow semantics are understood
4. Double-check dtype spelling: jnp.uint32 vs jnp.int32

Example fix

// before
out = lax.dot_general(a_i32, b_i32, ..., preferred_element_type=jnp.uint32)

// after
out = lax.dot_general(a_i32, b_i32, ..., preferred_element_type=jnp.int64)

import numpy as np
assert not (np.issubdtype(input_dtype, np.signedinteger) and
           not np.issubdtype(pref, np.signedinteger)), 'signedness mismatch'

def same_signedness(input_dtype, pref) -> bool:
    return np.issubdtype(input_dtype, np.signedinteger) == np.issubdtype(pref, np.signedinteger) or not np.issubdtype(input_dtype, np.integer)

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

When it happens

Trigger: lax.dot_general with int32 (signed) inputs and preferred_element_type=jnp.uint32, or uint8 inputs with int16 preferred.

Common situations: Trying to use unsigned accumulators to gain range; guessing dtype names (uint32 vs int32); porting C-style unsigned accumulation habits.

Related errors

• Input type is incompatible with `preferred_element_type`. Th
• `preferred_element_type` must not be narrower than the origi
• {} does not accept dtype {}. Accepted dtypes are subtypes of
• {name} does not accept dtype {dtype_to_string(aval.dtype)}.
• {} does not accept dtype {} at position {}. Accepted dtypes

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