Repository navigation
ENH: lazy_while (and friends) #303
Description
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 valThis 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.
Activity
lucascolley commentedon May 20, 2025
cc @crusaderky
[-]lazy_while(and friends)[/-][+]ENH: `lazy_while` (and friends)[/+]crusaderky commentedon May 20, 2025
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_loopis a good idea. Two reasons for this:jax.lax.while_loop(cond_fun, body_fun, init_val)intrinsically relies on the ability to jitcond_funandbody_funto run at speed, which most backends don't have.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 commentedon May 20, 2025
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_applydoes), 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 commentedon May 20, 2025
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.allsince it's easier to make the algorithm keep going rather than masking everything.purepani commentedon May 20, 2025
Currently the implementation of
celthat 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.pyThe compiled scipy implementations for
ellipkandellipeare used, though, so in those cases the whole loop would be slower(celis 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 commentedon May 20, 2025
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 commentedon May 20, 2025
I'm thinking a pattern like this would be useful for these cases:
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 tolazy_whilefor everything else.NeilGirdhar commentedon May 26, 2025
I think having some kind of
xpx.while_loopis 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_cpuprimitive? If so, where can I find that?purepani commentedon May 26, 2025
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_cpufunction for that comment.NeilGirdhar commentedon May 26, 2025
I don't understand? That's slower than XLA, right? Have you tested your algorithm under Jax's JIT?
Okay! Just thought it looked useful 😄
purepani commentedon May 26, 2025
Sorry, typo. I meant cython. Point is that it's already compiled code, so it is probably not much faster, if at all. I haven't tested it, so I could be wrong, but I think it's fair to assume there won't be a significant performance difference either way.
34j commentedon Jul 13, 2025
If the number of loops is rather small, falling back to a Python for loop should not be a problem; given that there are only three types of loops in JAX, it should not be too many lines, so I hope this could be implemented here to allow fast JIT compilation.
lucascolley commentedon Aug 12, 2026
See scipy/scipy#24956 for another potential use-case here. Still not totally clear what the API ought to look like.