Best practice: scan vs fori_loop/while_loop

[AdrienCorenflos]

Hi,

I'm a bit confused so as to what syntax I should use to loop over some data. In practice virtually any loop can be coded using a scan-like syntax but does that mean it should?

The only difference I am aware of is the support for backprop in the scan case whilst fori_loop (or while_loop). Is there any other hidden one?

Just to fix ideas below is a toy example which computes the same result, but has different behaviours in terms of gradient implementation:

Thanks

Adrien

import jax.numpy as jnp
from jax import lax, jit
from jax.test_util import check_grads

@jit
def scan_fun(x, y):
    # A loop using scan, note how I access y within the body
    n = x.shape[0]

    def body(carry, x):
        curr, i = carry
        return (curr + x * y[i], i+1), None

    (res, _), _ = lax.scan(body, (0., 0), x)
    return res / n

@jit
def loop_fun(x, y):
    # A loop using fori_loop, note how I access both x and y within the body
    n = x.shape[0]

    def body(i, curr): 
        return curr + x[i] * y[i]

    return lax.fori_loop(0, n, body, 0.) / n

n = 50
m = 150
x = jnp.arange(n, dtype=jnp.float32)
y = jnp.arange(m, step=2., dtype=jnp.float32)

assert loop_fun(x, y) == scan_fun(x, y)

check_grads(scan_fun, (x, y), 1, eps=1e-2)  # All good
check_grads(loop_fun, (x, y), 1, eps=1e-2)  # raises ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop. Try using lax.scan instead.

[mattjj]

Thanks for the question!

Our slogan is, "always scan when you can!"

This doesn't apply to your example, but in general it's also a good idea to use what distinguishes scan from fori_loop when you can, i.e. the scanned-over inputs and outputs rather than the loop carry (since fori_loop only has the loop carry). When you use scanned-over inputs and outputs instead of using the loop carry, it lets JAX generate more efficient differentiation code. The reason is pretty straightforward: we need to save data from each loop iteration for the forward pass to be consumed on the backward pass. When that data is in the loop carry we basically have to snapshot the whole loop carry for each iteration, but when the data is in scanned-over inputs and outputs we know we only have to save a particular slice for each iteration.

That comment doesn't apply to the scan_fun example here, though. While you're doing the slicing manually rather than treating y as a scanned-over input, closed-over values like y are handled differently from loop carry values for autodiff purposes, so I believe things are still as efficient as possible.

Zooming back out to the big picture, the only reason not to use scan is if you can't, because the number of loop iterations is not a function of just literals or array shapes but is instead a function of runtime values. Canonical examples are fixed-point iterations or rejection samplers, where we want to iterate until some data-dependent threshold is achieved; in those cases, we're forced to use a while_loop.

(In fact, we'd even like to implement fori_loop in terms of scan under the hood when possible, though that isn't enabled right now for some boring reason I can't remember.)

WDYT?

[jecampagne]

In fact this is the case PR #2414, and after discussion #14643, the case fori_loop cannot be used with JIT/jacrev are certainly rare (ie. one uses static_arguments in the right place). Now if one finds a use case where one cannot use fori_loop it would be illustrative.

[mattjj]

One other thought: my previous comment focused on the autodiff advantages of using scan instead of other loop constructs. But there may be codegen and execution performance advantages too! Because scan provides more information to the XLA compiler (namely a loop with a guaranteed static trip count, since XLA programs are shape-specialized), the compiler can often do more optimizations. So scan is often faster too, even without autodiff! Seems like a good deal.

[AdrienCorenflos]

Thanks for the explanation Matt, that's crystal clear. One point you raise w.r.t. the performance though makes me wonder. I don't see how this can be true when you account for the second output (the accumulated one), which is useless in a number of applications. That being said, I suppose this would be discarded statically if not used in a jitted function? Adrien EDIT: I apologize for the originally horrendous format of the reply (seems like the email answers don't bode well with GitHub yet)

…

On Fri, 24 Jul 2020, 18:34 Matthew Johnson, ***@***.***> wrote: One other thought: my previous comment focused on the autodiff advantages of using scan instead of other loop constructs. But there may be performance advantages too! Because scan provides more information to the XLA compiler (namely a loop with a guaranteed static trip count, since XLA programs are shape-specialized), the compiler can often do more optimizations. So scan is often faster too, even without autodiff! Seems like a good deal. — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <https://github.com/google/jax/issues/3850#issuecomment-663598272>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AEYGFZY4WYFRCDINKILVILDR5GSZFANCNFSM4PGX43JA> .

[AdrienCorenflos]

As an FYI, I'm posting my real use case for the question: I am doing a linear interpolation on some data I know is sorted. Of course the function below is numba'd in "real life"

import numpy as np

def sorted_interp(x, xp, fp):
    """ Interpolates linearly between xp and fp BUT makes the assumption that x is sorted,
    this allows the complextity to be O(n+m) instead of O(m * log(n)) and leverages array contiguity
    (not like the simpler searchsorted implementation). Equivalent to np.interp(np.sort(x), xp, fp)
    Parameters
    ----------
    x: (M) array
        locations at which to interpolate
    xp: (N) array
        locations of the given function values
    fp: (N) array
        values of the function at the given locations
    Returns
    -------
    out: (M) array
        the resulting interpolated values
    """
    m = x.shape[0]
    n = xp.shape[0]

    x = np.atleast_1d(x)

    j = 0
    xp_0 = xp[0]
    fp_0 = fp[0]

    res = np.empty(x.shape, fp.dtype)

    for i in range(m):
        x_i = x[i]

        if x_i <= xp_0:
            res[i] = fp_0
            continue
        else:
            while True:
                xp_j = xp[j]
                fp_j = fp[j]

                next_xp_j = xp[j + 1]
                next_fp_j = fp[j + 1]

                if x_i > next_xp_j:
                    if j + 1 == n - 1:
                        res[i] = next_fp_j
                        break
                    else:
                        j += 1
                        continue
                else:
                    if xp_j == next_xp_j:
                        val = fp_j
                    else:
                        val = fp_j + (next_fp_j - fp_j) * (x_i - xp_j) / (next_xp_j - xp_j)
                    res[i] = val
                    break
    return res

Now that I think of it, once I've implemented it in JAX, I may just leverage my use case (long live open source) to provide you guys with a jax implementation of np.interp as I think it's still missing. It may even prove faster than the numpy version...

def inv(argsort, arr):
    """Computes the inverse of the permutation, but might be better 
    for autodiff to simply compute the inverse permutation and then index the array
    """
    inverse = np.empty_like(arr)
    for i, p in enumerate(argsort):
        inverse[p] = arr[i]
    return inverse

def interp(x, xp, fp):
    """Equivalent of np.interp that doesn't rely on searchsorted"""
    argsort = np.argsort(x)
    sorted_x= x[argsort]
    interpolated_f = sorted_interp(sorted_x, xp, fp)
    return inv(argsort, interpolated_f)

[jakevdp]

Side-note: a useful trick is that you can invert an argsort by argsorting the argsort 😁

def inv(argsort, arr):
  return arr[np.argsort(argsort)]

[AdrienCorenflos]

Yes but it's a costly trick when one can do it in linear time :)

[shoyer]

You can also calculate an inverse permutation in single call to array indexing:

def inv(argsort, arr):
    inverse = np.empty_like(arr)
    inverse[argsort] = arr
    return inverse

In JAX, you would write this as inverse = jnp.empty_like(arr).at[argsort].set(arr)

[AdrienCorenflos]

Just opened a new issue to discuss specifics and not spam the "scan vs loop" question: #3860

[mattjj]

I don't see how this can be true when you account for the second output (the accumulated one), which is useless in a number of applications. That being said, I suppose this would be discarded statically if not used in a jitted function?

Do you mean the scanned-over output? Not only will that be ignored if unused in a jitted function, but also it's optional; you can just scan a function with a None as its second output.

I think we covered this one, so I'm going to close this issue!

[shoyer]

Should we try converting this issue into a “Discussion” for posterity?

…

On Mon, Jul 27, 2020 at 8:51 PM Matthew Johnson ***@***.***> wrote: I don't see how this can be true when you account for the second output (the accumulated one), which is useless in a number of applications. That being said, I suppose this would be discarded statically if not used in a jitted function? Do you mean the scanned-over output? Not only will that be ignored if unused in a jitted function, but also it's optional; you can just scan a function with a None as its second output. I think we covered this one, so I'm going to close this issue! — You are receiving this because you commented. Reply to this email directly, view it on GitHub <https://github.com/google/jax/issues/3850#issuecomment-664760258>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVVLBWVUNC2I2P5VURTR5ZDKFANCNFSM4PGX43JA> .

[mattjj]

@shoyer you just want to exercise your discussion-moving GitHub powers!

SGTM!

[bionicles]

how could we bail out of scans early to make a differentiable while loop?

[mattjj]

Great question. I think this is a TODO item for us (or fearless contributors!). It'd have to be built into scan itself, which currently always iterates up to the max iteration:

jax/jax/lax/lax_control_flow.py

Lines 1272 to 1274 in 0ec1e25

	def cond_fun(vals):
	i, *_ = vals
	return i < length

That'd be part of the puzzle, though there are likely other things I'm forgetting...

[shoyer]

This would be really nice to have for iterative algorithms like cg. Our current implementation with while_loop requires using the implicit function theorem for gradients, but it would be nice to differentiate through some fixed number of cg steps as well, e.g., like what was done for optimizing initial conditions for Poisson solves in https://arxiv.org/abs/2007.00016.

[shoyer]

Implementation wise, the other thing we'd need to figure out is what the fill values on the resulting array from the forward scan should be. Potentially this could be an application for something like the experimental auto-masking transformation.

[bionicles]

This jitted OK, you can pass it a lambda function which invokes a Stax layer. I resisted Flax because it's OOP but I'll try that next

"adaptive computation time via optimization"
from jax.experimental.optimizers import sgd
from jax.lax import scan, cond
from jax import numpy as jnp
from jax import grad, jit


sum_abs = lambda x: jnp.sum(jnp.abs(x))

def minimize(
    fun, inputs, optimizer=sgd, schedule=1e-3, maxiter=4, gtol=1e-2,
):
    """Optimizes inputs to minimize output of energy_fn
    Args:
        fun: maps x -> energy
        inputs: initial inputs
        optimizer: any jax.experimental.optimzier
        schedule: learning rate schedule for optimizer
        maxiter: maximum number of optimizer steps
        gtol: cutoff for gradients to stop early
    """
    init_optimizer, update_optimizer, get_inputs = optimizer(schedule)
    energy_gradient = grad(lambda x: jnp.sum(fun(x)))
    opt_state = init_optimizer(inputs)
    gradient = energy_gradient(inputs)
    sum_abs_gradient = sum_abs(gradient)
    carry = (0, opt_state, gradient, sum_abs_gradient, inputs)

    @jit
    def step_once(carry):
        (step, opt_state, gradient, sum_abs_gradient, inputs) = carry
        opt_state = update_optimizer(step, gradient, opt_state)
        inputs = get_inputs(opt_state)
        gradient = energy_gradient(inputs)
        sum_abs_gradient = sum_abs(gradient)
        carry = (step + 1, opt_state, gradient, sum_abs_gradient, inputs)
        return carry, None

    def noop(carry):
        return carry, None

    @jit
    def body_fn(carry, x):
        sum_abs_gradient = carry[-2]
        return cond(sum_abs_gradient < gtol, carry, step_once, carry, noop)

    carry, _ = scan(body_fn, carry, None, length=maxiter)
    return carry