Repository: google/jax

Issue:
bug in shape-checking rules for the case of fori_loop
The following three loops should all give the same results but the lax.fori_loop throws an error.
Could be related to #1681 

import jax.numpy as np
from jax.lax import fori_loop as fori_loop_jax

def fori_loop_python(lower, upper, body_fun, init_val):
    val = init_val
    for i in range(lower, upper):
        val = body_fun(i, val)
    return val

def body_fun(i,arr):
    a,b = arr
    return a-b[i]*b[i+1], b

s = np.ones([10,10])
r = s.shape[1]-1

E = 0
for i in range(r):
    E -= s[:, i] * s[:, i + 1]
print(E)

E, _ = fori_loop_python(0, s.shape[0]-1, body_fun, (0,s))
print(E)

E, _ = fori_loop_jax(0, r, body_fun, (0,s))[0]
print(E)


Discussion:
Thanks for raising this!

I tweaked the error message to print with some more newlines, to make things line up:

```
TypeError: body_fun output and input must have identical types, got
(ShapedArray(int32[]), ShapedArray(int32[]), (ShapedArray(float32[10]), ShapedArray(float32[10,10])))
and
(ShapedArray(int32[]), ShapedArray(int32[]), (ShapedArray(int32[]), ShapedArray(float32[10,10]))).
```

As you can see, there's an `int32[]` input where we get a `float32[10]` output. The issue is that in the Python case you're relying on `0` getting rank- and dtype-promoted to a `float32[10]`, and that's not allowed for `lax.fori_loop`.

This works (notice I removed a `[0]` from the end):

```python
E, _ = fori_loop_jax(0, r, body_fun, (np.zeros(10),s))
```

We could probably be clearer in our documentation that, unlike the pure Python analogue, the input and output types have to be the same, not just consistent up to NumPy-style broadcasting/promotion.

Resolving patch:
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -112,6 +112,14 @@ def fori_loop(lower, upper, body_fun, init_val):
   Unlike that Python version, ``fori_loop`` is implemented in terms of a call to
   ``while_loop``. See the docstring for ``while_loop`` for more information.
 
+  Also unlike the Python analogue, the loop-carried value ``val`` must hold a
+  fixed shape and dtype across all iterations (and not just be consistent up to
+  NumPy rank/shape broadcasting and dtype promotion rules, for example). In
+  other words, the type ``a`` in the type signature above represents an array
+  with a fixed shape and dtype (or a nested tuple/list/dict container data
+  structure with a fixed structure and arrays with fixed shape and dtype at the
+  leaves).
+
   Args:
     lower: an integer representing the loop index lower bound (inclusive)
     upper: an integer representing the loop index upper bound (exclusive)
@@ -155,6 +163,14 @@ def while_loop(cond_fun, body_fun, init_val):
   for jit-compiled functions, since native Python loop constructs in an ``@jit``
   function are unrolled, leading to large XLA computations.
 
+  Also unlike the Python analogue, the loop-carried value ``val`` must hold a
+  fixed shape and dtype across all iterations (and not just be consistent up to
+  NumPy rank/shape broadcasting and dtype promotion rules, for example). In
+  other words, the type ``a`` in the type signature above represents an array
+  with a fixed shape and dtype (or a nested tuple/list/dict container data
+  structure with a fixed structure and arrays with fixed shape and dtype at the
+  leaves).
+
   Another difference from using Python-native loop constructs is that
   ``while_loop`` is not reverse-mode differentiable because XLA computations
   require static bounds on memory requirements.
@@ -455,6 +471,13 @@ def scan(f, init, xs):
   for jit-compiled functions, since native Python loop constructs in an ``@jit``
   function are unrolled, leading to large XLA computations.
 
+  Finally, the loop-carried value ``carry`` must hold a fixed shape and dtype
+  across all iterations (and not just be consistent up to NumPy rank/shape
+  broadcasting and dtype promotion rules, for example). In other words, the type
+  ``c`` in the type signature above represents an array with a fixed shape and
+  dtype (or a nested tuple/list/dict container data structure with a fixed
+  structure and arrays with fixed shape and dtype at the leaves).
+
   Args:
     f: a Python function to be scanned of type ``c -> a -> (c, b)``, meaning
       that ``f`` accepts two arguments where the first is a value of the loop
@@ -968,7 +991,7 @@ def _check_tree_and_avals(what, tree1, avals1, tree2, avals2):
     raise TypeError(msg.format(what, tree1, tree2))
   if not all(safe_map(typematch, avals1, avals2)):
     msg = ("{} must have identical types, "
-           "got {} and {}.")
+           "got\n{}\nand\n{}.")
     raise TypeError(msg.format(what, tree_unflatten(tree1, avals1),
                                tree_unflatten(tree2, avals2)))
 


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
@@ -202,8 +202,12 @@ def testWhileTypeErrors(self):
     with self.assertRaisesRegex(TypeError,
         re.escape("body_fun output and input must have same type structure, got PyTreeDef(tuple, [*,*]) and *.")):
       lax.while_loop(lambda c: True, lambda c: (1., 1.), 0.)
-    with self.assertRaisesRegex(TypeError,
-        re.escape("body_fun output and input must have identical types, got ShapedArray(bool[]) and ShapedArray(float32[]).")):
+    with self.assertRaisesRegex(
+        TypeError,
+        "body_fun output and input must have identical types, got\\n"
+        "ShapedArray\(bool\[\]\)\\n"
+        "and\\n"
+        "ShapedArray\(float32\[\]\)."):
       lax.while_loop(lambda c: True, lambda c: True, np.float32(0.))
 
   def testNestedWhileWithDynamicUpdateSlice(self):
@@ -526,8 +530,12 @@ def testCondTypeErrors(self):
         re.escape("true_fun and false_fun output must have same type structure, got * and PyTreeDef(tuple, [*,*]).")):
       lax.cond(True,
                1., lambda top: 1., 2., lambda fop: (2., 2.))
-    with self.assertRaisesRegex(TypeError,
-        re.escape("true_fun and false_fun output must have identical types, got ShapedArray(float32[1]) and ShapedArray(float32[]).")):
+    with self.assertRaisesRegex(
+        TypeError,
+        "true_fun and false_fun output must have identical types, got\n"
+        "ShapedArray\(float32\[1\]\)\n"
+        "and\n"
+        "ShapedArray\(float32\[\]\)."):
       lax.cond(True,
                1., lambda top: np.array([1.], np.float32),
                2., lambda fop: np.float32(1.))
@@ -903,9 +911,12 @@ def testScanTypeErrors(self):
     with self.assertRaisesRegex(TypeError,
         re.escape("scan carry output and input must have same type structure, got * and PyTreeDef(None, []).")):
       lax.scan(lambda c, x: (0, x), None, a)
-    with self.assertRaisesRegex(TypeError,
-        re.escape("scan carry output and input must have identical types, "
-                  "got ShapedArray(int32[]) and ShapedArray(float32[]).")):
+    with self.assertRaisesRegex(
+        TypeError,
+        "scan carry output and input must have identical types, got\n"
+        "ShapedArray\(int32\[\]\)\\n"
+        "and\\n"
+        "ShapedArray\(float32\[\]\)."):
       lax.scan(lambda c, x: (np.int32(0), x), np.float32(1.0), a)
     with self.assertRaisesRegex(TypeError,
         re.escape("scan carry output and input must have same type structure, got * and PyTreeDef(tuple, [*,*]).")):