Repository: google/jax
Issue: fori_loop with disable_jit fails when no actual iterations happen

In a trivial call of `fori_loop`, i.e. when `upper == lower`, a normal call works just fine,
but a call with `disable_jit` causes a `TypeError`.

Can be reproduced by 
```python
from jax import lax, disable_jit


def test():
    non_trivial = lax.fori_loop(lower=0, upper=3, body_fun=lambda x, i: x + i, init_val=1)
    print(f'non-trivial result: {non_trivial}')
    trivial = lax.fori_loop(lower=0, upper=0, body_fun=lambda x, i: x + i, init_val=1)
    print(f'trivial result: {trivial}')


test()
# non-trivial result: 4
# trivial result: 1

with disable_jit():
    test()
# non-trivial result: 4
# TypeError: tree_map() missing 1 required positional argument: 'tree'
```

This is because the [non-jit implementation of `lax.scan`](https://github.com/google/jax/blob/cd3d37f4d14bdce8501795c759be7a5b9f79e157/jax/_src/lax/control_flow.py#L1286)
calls `tree_multimap` to stack the outputs, but there aren't any outputs to stack.

I suggest to check for this special case (of a trivial array to scan over) in the non-jit branch of scan and directly return the initial carry and a trivial stacked_y.

Happy to make a PR if this is the way you want to do it.

Changed files:
tests/lax_control_flow_test.py, jax/_src/lax/control_flow.py

diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
index 38a195148..26de4c9b0 100644
--- a/jax/_src/lax/control_flow.py
+++ b/jax/_src/lax/control_flow.py
@@ -210,6 +210,10 @@ def fori_loop(lower, upper, body_fun, init_val):
     use_scan = False
 
   if use_scan:
+    if config.jax_disable_jit and upper_ == lower_:
+      # non-jit implementation of scan does not support length=0
+      return init_val
+
     (_, result), _ = scan(_fori_scan_body_fun(body_fun), (lower_, init_val),
                           None, length=upper_ - lower_)
   else:
@@ -1284,6 +1288,8 @@ def scan(f: Callable[[Carry, X], Tuple[Carry, Y]],
       length, = unique_lengths
 
   if config.jax_disable_jit:
+    if length == 0:
+      raise ValueError("zero-length scan is not supported in disable_jit() mode because the output type is unknown.")
     carry = init
     ys = []
     maybe_reversed = reversed if reverse else lambda x: x
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index a6aa9e1f7..4b31a00bc 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -542,6 +542,23 @@ class LaxControlFlowTest(jtu.JaxTestCase):
       self.assertAllClose(cfun(x, num), np.sum(x[:num]), check_dtypes=False)
       self.assertAllClose(cfun(x, num), np.sum(x[:num]), check_dtypes=False)
 
+  def testForiLoopIssue8152(self):
+    y = lax.fori_loop(lower=0, upper=0, body_fun=lambda x, i: x + i, init_val=1.)
+    self.assertAllClose(y, 1., check_dtypes=False)
+
+    # trivial fori_loop should work - even when jit is disabled
+    with jax.disable_jit():
+      y = lax.fori_loop(lower=0, upper=0, body_fun=lambda x, i: x + i, init_val=1.)
+    self.assertAllClose(y, 1., check_dtypes=False)
+
+    # scan with length 0 should work with jit, but raise an error without
+    def should_raise_wo_jit():
+      carry, out = lax.scan(lambda c, x: (c + x, x), 0., np.array([]))
+      return carry
+    self.assertAllClose(should_raise_wo_jit(), 0., check_dtypes=False)
+    with jax.disable_jit():
+      self.assertRaises(ValueError, should_raise_wo_jit)
+
   def testCond(self):
     def fun(x):
       if x < 3: