Nested loop in jax

[gerdm]

Prior to #10278, I would've written a nested for-loop in jax to accumulate all possible values of a function f: [X,Y,Z] -> [] as follows:

with loops.Scope() as s:
    s.res = 0.0
    for x in s.range(X):
        for y in s.range(Y):
            for z in s.range(Z):
                s.res = s.res + f(X_vals[x], Y_vals[y], Z_vals[z])

I understand that this API is being deprecated in favour of lax.fori_loop. However, it's not totally clear to me how to re-write the previous for-loop using lax.fori_loop. What's the best way of doing so?

Also, why is the loops module being deprecated? I've found it pretty convenient when having to write nested for-loops in jax.

Thanks!

[jakevdp]

Assuming your function f handles broadcasting semantics, you could do this:

result = f(Xvals[:, None, None], Yvals[None, :, None], Zvals[None, None, :]).sum((0, 1, 2))

If not, you could do something similar with vmap:

f_mapped = vmap(vmap(vmap(f, in_axes=(0, None, None)), in_axes=(None, 0, None)), in_axes=(None, None, 0))(Xvals, Yvals, Zvals).sum((0, 1, 2))

I suspect either of these would be much more performant than fori_loop.

But if you need fori_loop, you can express this by nesting three loop calls:

from jax import lax, vmap

x = jnp.arange(5)
y = jnp.arange(10)
z = jnp.arange(15)

def f(x, y, z):
  return x * y * z

print(sum(f(x[i], y[j], z[k])
          for i in range(len(x))
          for j in range(len(y))
          for k in range(len(z))))
# 47250

print(f(x[None, None, :], y[None, :, None], z[:, None, None]).sum())
# 47250

print(vmap(vmap(vmap(f, in_axes=(0, None, None)), in_axes=(None, 0, None)), in_axes=(None, None, 0))(x, y, z).sum())
# 47250

print(lax.fori_loop(0, len(x),
    lambda i, vx: vx + lax.fori_loop(0, len(y),
      lambda j, vy: vy + lax.fori_loop(0, len(z),
        lambda k, vz: vz + f(x[i], y[j], z[k]),
        0),
      0),
    0))
# 47250

[gerdm]

Thanks for your answer @jakevdp!

I was looking for a way to avoid a vmap since the cardinality of my arrays are too big to fit in memory when computingf[x, y, z] for all elements.

[jakevdp]

Makes sense. Under some circumstances, I believe the XLA compiler can avoid materialization of intermediate results in an aggregation like this, but sometimes loops are necessary. See my edited answer for an example.

[mattjj]

And if you want to generate a loop, you can write as a triply-nested scan, which should lead to more efficient autodiff than the loops module would get you here:

import jax
import jax.numpy as jnp

def product_loop_nest(f, *xs):
  # for xs = (x1, x2, ...), this evaluates the value of `s` after the loop nest
  # s = 0.0
  # for i in range(len(x1)):
  #   for j in range(len(x2)):
  #     ...
  #       s = s + f(x1[i], x2[j], ...)

  # base case takes n+1 arguments representing slices (s, x1[i], x2[j], ...)
  loop = lambda s, *x_is: s + f(*x_is)

  # induction case: scan over one argument, eliminating it
  def scan_one_more(loop, x):
    def new_loop(s, *xs):
      s, _ = jax.lax.scan(lambda s, x_i: (loop(s, *xs, x_i), None), s, x)
      return s
    return new_loop

  # compose
  for x in reversed(xs):
    loop = scan_one_more(loop, x)
  return loop(0)


out = product_loop_nest(lambda i, j, k: i * j * k,
                        jnp.arange(3), jnp.arange(4), jnp.arange(5))
print(out)

expected = (jnp.arange(3)[:, None, None] * jnp.arange(4)[None, :, None] *
            jnp.arange(5)[None, None, :]).sum()
print(expected)

[mattjj]

Using scan is more memory- and time-efficient than fori_loop under autodiff, since the backward-pass updates generated by indexed reads on the forward pass are sparse with scan but not with fori_loop.