def fori_loop(lower, upper, body_fn, init_val):
  """Loop from `lower` to `upper` running `body_fn` starting from `init_val`.

  The semantics of `fori_loop` is as follows::

    def fori_loop(lower, upper, body_fn, init_val):
      val = init_val
      for i in range(lower, upper):
        val = body_fn(i, val)
      return val

  Args:
    lower: an integer representing the loop index lower bound (inclusive)
    upper: an integer representing the loop index upper bound (exclusive)
    body_fn: function of type `(int, a) -> a`.
    init_val: initial loop carry value of type `a`.

  Returns:
    Loop value from the final iteration.
  """
  if 'fori_loop' in backend():
    return backend()['fori_loop'](lower, upper, body_fn, init_val)
  # Use scan otherwise.
  def scanned_fn(loop_carry, _):
    i, x = loop_carry
    return (i + 1, body_fn(i, x)), None
  (_, result), _ = scan(
      scanned_fn, (lower, init_val), None, length=upper - lower)
  return result
Loop from `lower` to `upper` running `body_fn` starting from `init_val`.

  The semantics of `fori_loop` is as follows::

    def fori_loop(lower, upper, body_fn, init_val):
      val = init_val
      for i in range(lower, upper):
        val = body_fn(i, val)
      return val

  Args:
    lower: an integer representing the loop index lower bound (inclusive)
    upper: an integer representing the loop index upper bound (exclusive)
    body_fn: function of type `(int, a) -> a`.
    init_val: initial loop carry value of type `a`.

  Returns:
    Loop value from the final iteration.
  
trax/fastmath/ops.py