jax-ml/jax · error · TypeError

mulhi operands must have the same dtype, got {dtype} and {y_

Error message

mulhi operands must have the same dtype, got {dtype} and {y_dtype}

What it means

mulhi requires both operands to share the same integer dtype; mixed dtypes (e.g. int32 and int64) raise TypeError.

Source

Thrown at jax/_src/lax_reference.py:120

add = np.add
sub = np.subtract

def mul(x, y, /, *, out_dtype=None):
  if out_dtype is not None:
    x = np.astype(x, out_dtype)
    y = np.astype(y, out_dtype)
  return np.multiply(x, y)


def mulhi(x, y):
  x = np.asarray(x)
  y = np.asarray(y)
  dtype = x.dtype
  if not np.issubdtype(dtype, np.integer):
    raise TypeError(f'mulhi requires integer inputs, got {dtype}')
  if dtype != y.dtype:
    raise TypeError(
        f'mulhi operands must have the same dtype, got {dtype} and {y.dtype}'
    )
  info = np.iinfo(dtype)
  bits = info.bits
  is_signed = np.issubdtype(dtype, np.signedinteger)
  # For 64-bit inputs, use Python object dtype for arbitrary precision.
  if bits == 64:
    widen_dtype = np.dtype(object)
  else:
    widen_bits = bits * 2
    widen_dtype = np.dtype(f'{"i" if is_signed else "u"}{widen_bits // 8}')
  prod = x.astype(widen_dtype) * y.astype(widen_dtype)
  return (prod >> bits).astype(dtype)


def div(lhs, rhs):
  if dtypes.issubdtype(dtypes.result_type(lhs), np.integer):
    quotient = np.floor_divide(lhs, rhs)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast both operands to the same dtype before calling: np.int32 promotion of both
2. Use jnp int dtype explicitly on all inputs (np.arange(..., dtype=np.int32))

Example fix

# before
lax_reference.mulhi(a.astype(np.int32), b)  # b is int64
# after
lax_reference.mulhi(a.astype(np.int32), b.astype(np.int32))

x, y = np.asarray(x), np.asarray(y)
assert x.dtype == y.dtype, f'{x.dtype} vs {y.dtype}'

When it happens

Trigger: lax_reference.mulhi(np.int32(x), np.int64(y)) or mixing Python ints (int64) with explicit int32 arrays.

Common situations: Fixed-point arithmetic where one operand came from np.arange defaults and the other was explicitly cast; NumPy scalar promotion differences across versions.

Related errors

• mulhi requires integer inputs, got {dtype}
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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