jax-ml/jax · error · TypeError

lower and upper arguments to fori_loop must have equal types

Error message

lower and upper arguments to fori_loop must have equal types, got {} and {}

What it means

fori_loop requires lower and upper bounds to have dtypes that can be unified: equal dtypes, or one side a weakly-typed integer that can adopt the other's dtype. If both are concrete arrays with different, non-weak dtypes (e.g. int32 vs int64, or int32 vs float32), JAX raises this TypeError because loop bounds must have a single consistent type.

Source

Thrown at jax/_src/lax/control_flow/loops.py:2623

  upper_dtype = lax.dtype(upper)
  if lower_dtype == upper_dtype:
    dtype = lower_dtype
  else:
    # As a special case: allow promotion of weak integers (e.g., Python scalars)
    # This improves the ergonomics if one but not both of the loop bounds is a
    # scalar.
    dtype = None
    if (np.issubdtype(lower_dtype, np.signedinteger) and
        np.issubdtype(upper_dtype, np.signedinteger)):
      lower_weak = dtypes.is_weakly_typed(lower)
      upper_weak = dtypes.is_weakly_typed(upper)
      if lower_weak and not upper_weak:
        dtype = upper_dtype
      elif not lower_weak and upper_weak:
        dtype = lower_dtype

    if dtype is None:
      raise TypeError("lower and upper arguments to fori_loop must have equal "
                      f"types, got {lower_dtype.name} and {upper_dtype.name}")

  # If we can specialize on the trip count, call scan instead of a while_loop
  # to enable efficient reverse-mode differentiation.
  lower_ = upper_ = 0
  if core.is_concrete(lower) and core.is_concrete(upper):
    try:
      lower_ = int(lower)
      upper_ = int(upper)
    except (TypeError, core.InconclusiveDimensionOperation):
      use_scan = False
    else:
      use_scan = True
  else:
    use_scan = False

  body_fun_dbg = api_util.debug_info("fori_loop", body_fun,
                                     (0, init_val), {})

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast both bounds to the same dtype explicitly: lax.fori_loop(lax.convert_element_type(lower, jnp.int64), ... ) or pass Python ints which are weakly typed
2. Use plain Python ints for static bounds so weak typing lets JAX unify them
3. Trace where the mismatched dtypes originate (np.array defaults differ by platform) and normalize at the source, e.g. np.int64(0)

Example fix

// before
lo, hi = np.int32(0), np.int64(n)
lax.fori_loop(lo, hi, body, init)
// after
lo, hi = int(0), int(n)
lax.fori_loop(lo, hi, body, init)

import numpy as np, jax.numpy as jnp
def bounds_compatible(lower, upper):
    ld, ud = jnp.result_type(lower), jnp.result_type(upper)
    return ld == ud or (np.issubdtype(ld, np.integer) != np.issubdtype(ud, np.integer)) is False and (ld.weak or ud.weak)

def same_dtype_or_weak(lo, hi) -> bool:
    lt, ut = jax.dtype(lo), jax.dtype(hi)
    return lt == ut or lt.weak or ut.weak

When it happens

Trigger: Calling lax.fori_loop(lower, upper, ...) where lower is e.g. np.int32(0) and upper is np.int64(n), or one bound is an int32 array and the other a Python value promoted to float32; only when neither is a weakly-typed scalar that can be safely adopted.

Common situations: Mixing numpy scalars of different widths (np.int32 vs np.int64 common on Windows vs Linux), computing bounds from differently-typed arrays (index arrays vs shape-derived values), or passing one bound as jnp.float32 and the other as int.

Related errors

• lax.fori_loop: body_fun argument should be callable.
• iteration over a 0-d array
• Incorrect output dtype for return value #{i}: Expected: {out
• Formatting arguments to checkify.check need to be PyTrees of
• check_error takes an Error as argument, got type {type(error

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