jax-ml/jax · error · TypeError

{name} argument type error: {lhs.dtype}, {rhs.dtype}

Error message

{name} argument type error: {lhs.dtype}, {rhs.dtype}

What it means

Raised in the dot_general dtype rule when lhs and rhs have equal 'type properties' (precision class) but different dtypes, e.g. two different-width floats or mismatched signed/unsigned ints of the same rank. JAX refuses to silently pick a winner between same-class types.

Source

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

  del dimension_numbers  # unused
  # We're mostly matching XLA's logic here, namely in shape_inference.cc and
  # primitive_util.h's HigherPrecisionType, e.g.
  # https://github.com/openxla/xla/blob/ea3a841768d0dcf192e5820c9b25c34c73f2226a/xla/primitive_util.h#L329
  def type_properties(dt):
    c = _real_dtype(dt) if dtypes.issubdtype(dt, np.complexfloating) else dt
    return (dtypes.issubdtype(dt, np.complexfloating),
            dtypes.finfo(c).maxexp if dtypes.issubdtype(c, np.floating) else -1,
            dtypes.finfo(c).nmant  if dtypes.issubdtype(c, np.floating) else -1,
            _bit_width(c),
            not dtypes.issubdtype(c, np.unsignedinteger))
  lhs_prop, rhs_prop = type_properties(lhs.dtype), type_properties(rhs.dtype)
  if lhs_prop > rhs_prop:
    result_dtype = lhs.dtype
  elif rhs_prop > lhs_prop:
    result_dtype = rhs.dtype
  else:
    if lhs.dtype != rhs.dtype:
      raise TypeError(f'{name} argument type error: {lhs.dtype}, {rhs.dtype}')
    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)):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Explicitly cast both operands to a common dtype with .astype() before the dot
2. Use preferred_element_type after unifying input dtypes if you only want output widening
3. Audit upstream data loading for dtype drift (e.g. dataset float64 vs model float32)

Example fix

# before
out = lax.dot_general(x, w, (((1,), (0,)), ((), ())))  # x f32, w f64
# after
out = lax.dot_general(x.astype(w.dtype), w, (((1,), (0,)), ((), ())))

assert lhs.dtype == rhs.dtype or lhs.dtype.kind != rhs.dtype.kind, \
    f'mixed dtypes: {lhs.dtype} vs {rhs.dtype}'
# stronger: promote explicitly
lhs, rhs = jnp.promote_types(lhs, rhs), jnp.promote_types(lhs, rhs) if False else (lhs.astype(lhs.dtype), rhs)
# simplest correct check:
assert lhs.dtype == rhs.dtype

def same_dtype(lhs, rhs):
    return lhs.dtype == rhs.dtype

When it happens

Trigger: jax.lax.dot_general (or matmul paths hitting it) with operands like f32 @ f64, i32 @ u32, or bf16 @ f16 — same precision class, different dtypes.

Common situations: Mixing a model stored in float64 with float32 inputs; mixed-precision attention where q/k dtypes differ; implicit numpy scalar promotion expectations not applying in lax.

Related errors

• Unsupported {preferred_element_type=}
• linearized function called on tangent values inconsistent wi
• Incorrect output dtype for return value #{i}: Expected: {out
• {} function carry input and carry output must have equal typ
• lower and upper arguments to fori_loop must have equal types

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