Confused on fori_loop usage

[bhchiang]

I'm trying to use fori_loop in a simple example, but running into a bunch of errors:

import jax
import jax.numpy as jnp

init_val = (jnp.ones(2), jnp.zeros(2))


def f(a):
    return a


jax.lax.fori_loop(1, 5, f, init_val)

This gives me TypeError: f() takes 1 positional argument but 2 were given, although the documentation says "init_val – initial loop carry value of type a.", implying a single input argument (here a tuple).

I changed it to take in two input arguments.

init_val = (jnp.ones(2), jnp.zeros(2))


def f(a, b):
    return a, b


jax.lax.fori_loop(1, 5, f, init_val)

I get this error:

TypeError: body_fun output and input must have same type structure, got PyTreeDef(tuple, [*,*,PyTreeDef(tuple, [*,PyTreeDef(tuple, [*,*])])]) and PyTreeDef(tuple, [*,*,PyTreeDef(tuple, [*,*])]).

Changing it to return a tuple doesn't work either, nor does passing in init_val as (init_val,).

init_val = (jnp.ones(2), jnp.zeros(2))


def f(a, b):
    return (a, b)


jax.lax.fori_loop(1, 5, f, init_val)

Extremely confused on how to properly use this function. I'm not sure how to interpret the body_fun type structure output in the error above. When I print the tree structure for init_val, I just have the following as expected.

[DeviceArray([1., 1.], dtype=float32), DeviceArray([0., 0.], dtype=float32)] PyTreeDef(tuple, [*,*])

Any help appreciated, thanks.

[mattjj]

Great points. We should improve these error messages. (Actually this thread might work better as an Issue rather than a Discussion... not sure if we can move it somehow.)

I think the key is that fori_loop passes in a loop counter variable, but doesn't expect it to be passed back out. Here's an example:

import jax
import jax.numpy as jnp

init_val = (jnp.ones(2), jnp.zeros(2))

def f(i, a):
    return a

jax.lax.fori_loop(1, 5, f, init_val)

This snippet from the docstring is meant to convey the loop counter idea (notice how i is passed into body_fun but not returned by it):

The semantics of ``fori_loop`` are given by this Python implementation::

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

Do you have suggestions for how we might improve the docstring? Maybe if we give an example of calling fori_loop?