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.


Discussion:
Thanks for the report - it looks like you've diagnosed the issue. Would you like to prepare a PR? If you're not able to at the moment, I'm happy to look into a fix. Thanks!

Resolving patch:
diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
--- 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, init, xs, length=None):
       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


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -542,6 +542,23 @@ def body_fun(i, state):
       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: