Converting a lax.while_loop to lax.scan

[mjhoover1]

Hi, I have created a lax.while_loop which is functional with a loop guard which depends on a float and <= such as this:

step = 0
while .1 * step <= 10.:
 # Do something
 step += 1

Here it is:

def body(carry):
   # Do something
   return carry + 1

cond_fun = lambda carry: .1 * carry <= 10.

lax.while_loop(cond_fun, body, 0)

In an attempt to make it differentiable, is it possible to also convert this code to use lax.scan?

[jakevdp]

The easiest way to convert it to a scan would probably be to use lax.fori_loop, which is implemented via scan when the start and end ranges are concrete. For your function, it might look like this:

from jax import lax
lax.fori_loop(0, 101, lambda i, x: x + 1, 0)

In general, scan is only applicable when you know a priori how many loop iterations will be required, so it may not be possible depending on what your actual use-case is.

[jakevdp]

That tells me that your start and end indices are not concrete, so the fori_loop is lowering to a while_loop rather than a scan. If you have a concrete range, it lowers to scan and reverse-mode autodiff is not a problem:

from jax import lax, grad
def f(x):
  return lax.fori_loop(0, 101, lambda i, x: x + 1, x)

print(grad(f)(0.0))
# DeviceArray(1., dtype=float32, weak_type=True)

If your ranges are data-dependent, then scan is not possible (the number of repetitions must be known at compile-time) and while_loop with forward-mode autodiff is the only option.

[mjhoover1]

I see, and appreciate the clarification! Although, with regards to attempting forward-mode auto diff is it possible to use value_and_grad. I don't see anything in the documentation which performs a forward-mode auto diff except jax.jacfwd which seems to calculate the Jacobian. However, the output is just a scalar, so it wouldn't be necessary to calculate the Jacobian.

[mjhoover1]

I believe the solution is using jvp, but value_and_grad doesn't have this as a feature parameter to decide the differentiation used which would be nice. I'll do some digging to see if I can get it working. Thanks.

[jakevdp]

no, there's no way to use grad or value_and_grad in forward mode, unfortunately. But the jacobian of a scalar-valued function is essentially what grad gives you, so you may be able to use jacfwd directly.

[mjhoover1]

Hats off to you @jakevdp for being a very responsive and detailed maintainer of JAX!