jax-ml/jax · error · TypeError

`preferred_element_type` must not be narrower than the origi

Error message

`preferred_element_type` must not be narrower than the original type, got preferred_element_type of {preferred_element_type} for result type of {result_dtype}.

What it means

Raised by _maybe_upcast (replicating XLA's MaybeUpcast) when preferred_element_type has a smaller bit width than the natural result type and the result is not floating — e.g. asking for int8 output from int32 inputs. Non-float results cannot be safely narrowed.

Source

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

    result_dtype = lhs.dtype
  has_algorithm = isinstance(precision, (DotAlgorithm, DotAlgorithmPreset))
  return _maybe_upcast(result_dtype, preferred_element_type,
                       check_bit_width=not has_algorithm)

def _bit_width(d):
  if dtypes.issubdtype(d, np.inexact): return dtypes.finfo(d).bits
  elif dtypes.issubdtype(d, np.integer): return dtypes.iinfo(d).bits
  elif d == np.dtype('bool'): return 1
  else: assert False, d  # should be unreachable, open an issue!

def _maybe_upcast(result_dtype, preferred_element_type, check_bit_width):
  # replicates the logic in shape_inference.cc's MaybeUpcast
  if (preferred_element_type is None or
      result_dtype == preferred_element_type):
    return result_dtype
  if (check_bit_width and not dtypes.issubdtype(result_dtype, np.floating) and
      _bit_width(preferred_element_type) < _bit_width(result_dtype)):
    raise TypeError("`preferred_element_type` must not be narrower than the "
                    "original type, got preferred_element_type of "
                    f"{preferred_element_type} for result type of "
                    f"{result_dtype}.")
  return preferred_element_type

def _dot_general_transpose_lhs(g, x, y, *, dimension_numbers, precision,
                               preferred_element_type: DTypeLike | None,
                               out_sharding, swap_ans=False):
  (x_contract, y_contract), (x_batch, y_batch) = dimension_numbers
  x_ndim = x.aval.ndim
  x_kept = remaining(range(x_ndim), x_contract, x_batch)
  y_kept = remaining(range(np.ndim(y)), y_contract, y_batch)
  if swap_ans:
    ans_batch, ans_y, _ = ranges_like(x_batch, y_kept, x_kept)
  else:
    ans_batch, _, ans_y = ranges_like(x_batch, x_kept, y_kept)
  dims = ((ans_y, y_kept), (ans_batch, y_batch))
  x_contract_sorted_by_y = list(np.take(x_contract, np.argsort(y_contract)))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Choose preferred_element_type at least as wide as the input/result integer type
Cast inputs down first (x.astype(jnp.int8)) instead of narrowing the output
With DotAlgorithm, specify input/output types in the algorithm instead

Example fix

# before
out = lax.dot_general(a, b, dn, preferred_element_type=jnp.int8)  # a,b int32
# after
out = lax.dot_general(a.astype(jnp.int8), b.astype(jnp.int8), dn,
                      preferred_element_type=jnp.int32)

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
if preferred_element_type is not None:
    w_in = np.dtype(preferred_element_type).itemsize * 8
    w_out = np.dtype(result_dtype).itemsize * 8
    assert result_dtype.kind == 'f' or w_in >= w_out

Type guard

def upcast_ok(result_dtype, pet):
    import numpy as np
    return pet is None or result_dtype == pet or (result_dtype.kind == 'f' or np.dtype(pet).itemsize >= np.dtype(result_dtype).itemsize)

Prevention

Remember preferred_element_type sets accumulator/output type, not a cast
Cast inputs explicitly for quantized paths

When it happens

Trigger: lax.dot_general(..., preferred_element_type=jnp.int8) with integer operands of wider dtype; also paths without a DotAlgorithm (check_bit_width True).

Common situations: Quantization-style code trying to force a narrow accumulator type; confusing preferred_element_type (output/accumulator type) with input casting.

Related errors

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