Converting fori_loop to scan

[maxwass]

What I have below is just a fancy for-loop over step() which loops ~1e5 times and returns 2 jnp.ndarrays. I've implemented this in jax.lax.fori_loop for significant speed-up in compilation time, but this seems to lack Reverse-mode differentiation abilities, and thus I want to convert it to use jax.lax.scan.

def step(x: jnp.ndarray,
         w_prev: jnp.ndarray,
         lam_prev: jnp.ndarray,
         alpha: float,
         beta: float,
         D: jnp.ndarray
         ):
    # pure function: mostly jnp matmuls, dots, etc
    ...
    return w, lam

from functools import partial
@jit
def dpg_steps(x: jnp.ndarray,
              w_init: jnp.ndarray,
              lam_init: jnp.ndarray,
              alpha: float,
              beta: float,
              num_steps: int,
              D: jnp.ndarray
              ):
    f = partial(step, x=x, alpha=alpha, beta=beta, D=D)

    def body_fun(i, carry):
        w, lam = carry
        w, lam = f(w_prev=w, lam_prev=lam)
        return w, lam

    w, lam = jax.lax.fori_loop(0, num_steps, body_fun, (w_init, lam_init))
    return w, lam

I've tested this and it works (compiles and gives correct outputs). But when I used it to compute gradients w.r.t alpha and beta I get:

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop. Try using lax.scan instead.

So I am wondering if someone could guide me on the best way to convert this code to use scan?

[jakevdp]

fori_loop is implemented in terms of scan when it can be – and "can be" is controlled by whether the number of iterations is static. In your case, I think all you need in order for your function to be reverse-mode differentiable is to mark num_steps as a static parameter:

@partial(jit, static_argnames=['num_steps'])
def dpg_steps(...

[maxwass]

It seems to work, thanks!