ENH: `lazy_while` (and friends)

[purepani]

While converting to the array api for magpylib in magpylib/magpylib#844, a few of the functions(algorithms for some elliptic integrals) have the unavoidable pattern

tol = 1e-6
while (xp.all(val < tol)):
    val = make_val_better(val)
return val

This doesn't directly work for jit in jax(and I assume other lazy backends). jax offers a jax.lax.while_loop primative to replace this, so it would be good to dispatch to that with an array_api_extra.lazy_while.

There's also a jax.lax.fori_loop primitive that may be similarly useful. I have never used dask, and have no clue how to do this in dask, but I assume it also needs to be dispatched somehow.

I'll likely do this as part of magpylib/magpylib#844, and I can upstream it after.

[lucascolley]

cc @crusaderky

[crusaderky]

I'm aware that the lack of support for converging algorithms in the Array API is a major blind spot.
However I'm not entirely sure that writing a shim around jax.lax.while_loop is a good idea. Two reasons for this:

• Backends that don't support it natively (all others AFAIK?), which thus need to fall back to a pure python loop, will be extremely slow if the number of iterations is large-ish. Not sure how common is the use case for a small-ish number of iterations, with each individual iteration consuming a substantial amount of time?
• The syntax of jax.lax.while_loop(cond_fun, body_fun, init_val) intrinsically relies on the ability to jit cond_fun and body_fun to run at speed, which most backends don't have.

I have never used dask, and have no clue how to do this in dask, but I assume it also needs to be dispatched somehow.

Dask is notoriously bad at running converging algorithms.
It's currently impossible to define a converging algorithm in pure dask; writing one in dask.distributed requires spawning a nanny task, which is fragile to worker failures and doesn't allow for chunking without very complicated hoops.
Here's an old proposal to fix it, never actioned: dask/distributed#5671

[crusaderky]

It's currently impossible to define a converging algorithm in pure dask

To be more precise: it's impossible to write a converging algorithm with an exit condition that relies on a global reduction, without squashing all chunks together (like lazy_apply does), which is extremely bad for performance.
Elementwise exit conditions (end iteration on a chunk when all points of a chunk individually satisfy a condition; different chunks can take different number of iterations) are simpler and can be implemented on top of da.map_blocks. But that requires information about the algorithm being iterated upon.

[purepani]

Backends that don't support it natively (all others AFAIK?), which thus need to fall back to a pure python loop, will be extremely slow if the number of iterations is large-ish. Not sure how common the use case is for a small-ish number of iterations, with each individual iteration consuming a substantial amount of time?

At least for the use case of cel(which can give all the elliptic integrals), the algorithm only requires somewhere between 3-5 iterations, so dispatching to a slower while loop generally isn't a big deal in that case, though I do understand it not being a great idea for a generalized use case.

This specific algorithm is also an element wise exit condition, so that's not a huge deal for this case. I had just done xp.all since it's easier to make the algorithm keep going rather than masking everything.

[purepani]

Currently the implementation of cel that is used in the library is already a while loop in python: https://github.com/magpylib/magpylib/blob/main/src%2Fmagpylib%2F_src%2Ffields%2Fspecial_cel.py

The compiled scipy implementations for ellipk and ellipe are used, though, so in those cases the whole loop would be slower(cel is a generalized version of those functions).

This is of course all specific to this algorithm, and the points you make are certainly important for more generalized use.

[purepani]

I haven't looked so closely at the other algorithms, but wouldn't be surprised if all these elliptic integrals would benefit from being able to do something like this:
https://docs.scipy.org/doc/scipy/reference/special.html#elliptic-functions-and-integrals

They would severely degrade the performance for numpy, though, so it would still be a good idea to dispatch to the compiled version when possible in these cases.

[purepani]

I'm thinking a pattern like this would be useful for these cases:

def f_iterate_step(x):...
if is_cpu(x):
    val=xpx.lazy_apply(compiled_f, x, use_numpy=True)
else:
    val=xpx.lazy_while(f_iterate_step, x, lambda x: x<tol)

Edit: I guess it's not exactly like this since we probably want Jax on the CPU to be under lazy_while, but the idea is the same: dispatch to a numpy compiled version under some conditions and dispatch to lazy_while for everything else.

[NeilGirdhar]

I think having some kind of xpx.while_loop is a good idea. However:

You may want to consider a more general while loop like the one shown here. In particular, it has three "kinds": lax, checkpointed, and bounded. These will affect the differentiability of your algorithm, among other things.

Even under the CPU, Jax may be a lot faster than Numpy after the Jit runs, right?

If you don't call the Jax code, you'll lose differentiability, which may matter to some people using xpx.

Is there already an is_cpu primitive? If so, where can I find that?

[purepani]

Even under the CPU, Jax may be a lot faster than Numpy after the Jit runs, right?

In my case, not necessarily since the numpy code for the algorithm is currently compiled with cpython; I was focusing on my specific case for that comment, but yeah, for general situations, this wouldn't be true.

I made up the is_cpu function for that comment.

[NeilGirdhar]

In my case, not necessarily since the numpy code for the algorithm is currently compiled with cpython;

I don't understand? That's slower than XLA, right? Have you tested your algorithm under Jax's JIT?

I made up the is_cpu function for that comment.

Okay! Just thought it looked useful 😄