{"record":{"id":"12bb9fd64c552da5","repo":"jax-ml/jax","slug":"lower-and-upper-arguments-to-fori-loop-must-have-e","errorCode":null,"errorMessage":"lower and upper arguments to fori_loop must have equal types, got {} and {}","messagePattern":"lower and upper arguments to fori_loop must have equal types, got (.+?) and (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2623,"sourceCode":"  upper_dtype = lax.dtype(upper)\n  if lower_dtype == upper_dtype:\n    dtype = lower_dtype\n  else:\n    # As a special case: allow promotion of weak integers (e.g., Python scalars)\n    # This improves the ergonomics if one but not both of the loop bounds is a\n    # scalar.\n    dtype = None\n    if (np.issubdtype(lower_dtype, np.signedinteger) and\n        np.issubdtype(upper_dtype, np.signedinteger)):\n      lower_weak = dtypes.is_weakly_typed(lower)\n      upper_weak = dtypes.is_weakly_typed(upper)\n      if lower_weak and not upper_weak:\n        dtype = upper_dtype\n      elif not lower_weak and upper_weak:\n        dtype = lower_dtype\n\n    if dtype is None:\n      raise TypeError(\"lower and upper arguments to fori_loop must have equal \"\n                      f\"types, got {lower_dtype.name} and {upper_dtype.name}\")\n\n  # If we can specialize on the trip count, call scan instead of a while_loop\n  # to enable efficient reverse-mode differentiation.\n  lower_ = upper_ = 0\n  if core.is_concrete(lower) and core.is_concrete(upper):\n    try:\n      lower_ = int(lower)\n      upper_ = int(upper)\n    except (TypeError, core.InconclusiveDimensionOperation):\n      use_scan = False\n    else:\n      use_scan = True\n  else:\n    use_scan = False\n\n  body_fun_dbg = api_util.debug_info(\"fori_loop\", body_fun,\n                                     (0, init_val), {})","sourceCodeStart":2605,"sourceCodeEnd":2641,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2605-L2641","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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","Use plain Python ints for static bounds so weak typing lets JAX unify them","Trace where the mismatched dtypes originate (np.array defaults differ by platform) and normalize at the source, e.g. np.int64(0)"],"exampleFix":"// before\nlo, hi = np.int32(0), np.int64(n)\nlax.fori_loop(lo, hi, body, init)\n// after\nlo, hi = int(0), int(n)\nlax.fori_loop(lo, hi, body, init)","handlingStrategy":"validation","validationCode":"import numpy as np, jax.numpy as jnp\ndef bounds_compatible(lower, upper):\n    ld, ud = jnp.result_type(lower), jnp.result_type(upper)\n    return ld == ud or (np.issubdtype(ld, np.integer) != np.issubdtype(ud, np.integer)) is False and (ld.weak or ud.weak)","typeGuard":"def same_dtype_or_weak(lo, hi) -> bool:\n    lt, ut = jax.dtype(lo), jax.dtype(hi)\n    return lt == ut or lt.weak or ut.weak","tryCatchPattern":null,"preventionTips":["Use Python ints for loop bounds whenever possible","Normalize bound scalars to one dtype at creation: np.int64 / jnp.int32 consistently","Watch for cross-platform numpy default int width differences (Windows int32)"],"tags":["jax","fori-loop","dtype-mismatch","typeerror"],"backgroundTag":"dtype-mismatch-in-arguments","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}