jax-ml/jax · error · TypeError

`preferred_element_type` must not be narrower than the origi

Q: What causes "`preferred_element_type` must not be narrower than the original type." in jax-ml/jax?

Use a preferred type with itemsize >= input itemsize (e.g. float32 or float64 for float32 inputs)

Error message

`preferred_element_type` must not be narrower than the original type.

What it means

dot_general's preferred_element_type (an accumulation/intermediate precision) must not be narrower in bits than the input dtype — you cannot request int8 accumulation for int16 inputs or float16 accumulation for float32 inputs, since that would lose precision silently. TypeError is raised when itemsize(preferred) < itemsize(input).

Source

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

    # 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 "
           f"lhs_batch of {lhs_batch} and lhs_contracting of {lhs_contracting} "
           f"for lhs of rank {lhs.ndim}")
    raise TypeError(msg)
  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, rhs.ndim))
             for d in (rhs_contracting, rhs_batch)):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use a preferred type with itemsize >= input itemsize (e.g. float32 or float64 for float32 inputs)
If lower precision is intended, downcast the inputs themselves first, then dot at that precision
For bf16 speedups, cast inputs to bfloat16 and use matching preferred type
Leave preferred_element_type=None to use input precision

Example fix

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

# after
a, b = a_f32.astype(jnp.bfloat16), b_f32.astype(jnp.bfloat16)
out = lax.dot_general(a, b, ..., preferred_element_type=jnp.bfloat16)

Defensive patterns

Strategy: validation

Validate before calling

if preferred is not None:
    assert np.dtype(preferred).itemsize >= np.dtype(input_dtype).itemsize, 'too narrow'

Type guard

def wide_enough(input_dtype, preferred) -> bool:
    return np.dtype(preferred).itemsize >= np.dtype(input_dtype).itemsize

Try / catch

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

Prevention

Never request accumulation narrower than the inputs
Downcast inputs first if lower precision is the goal
Centralize precision policy in one helper that validates itemsizes

When it happens

Trigger: lax.dot_general(a_float32, b_float32, ..., preferred_element_type=jnp.float16); int16 inputs with int8 preferred; also triggered via jnp.dot paths that forward precision settings on TPUs.

Common situations: Trying to speed up matmuls by requesting lower-precision accumulation; TF/TPU precision presets (bfloat16 semantics) applied to float32 inputs incorrectly; confusing preferred_element_type with output dtype.

Related errors

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