jax-ml/jax · error · TypeError

mulhi operands must have the same dtype, got

Error message

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

What it means

mulhi's high-word extraction assumes both operands share the same integer width. The reference implementation fires when x.dtype differs from y.dtype, e.g. multiplying int32 by int64, which would change the number of high bits.

Solutions

Cast both operands to the same dtype before calling: np.int32 promotion of both
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))

Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

Set dtype= on all producing ops (arange, ones) to a fixed int width

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

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

Appendix: 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)