[FR] support batching rule for `while`

Currently, vmap seems do not support a function which involves lax.while_loop. For example, the following script will throw NotImplementedError: Batching rule for 'while' not implemented:

import jax.numpy as np
from jax import jit, lax, vmap

@jit
def f(x):
    return lax.while_loop(lambda x: x <= 0, lambda x: x + 1, x)

g = vmap(f)
y = g(np.arange(3.))

In case it is complicated to support while_loop for vmap, then could someone let me know other alternatives? I know some other alternatives from gufuncs tutorial such as using onp.vectorize or python for loop but they are slow. To make g fast, can I write f in C/Cython code and use vmap with it? Or should I write the whole g in C/Cython code? I have little experience with C/Cython but will try to learn if that is the only option.

I just want to use g as a primitive function and will define jvp_rule for it separately.

Thank you for any help in advance!

[mattjj]

@jonasrauber @fehiepsi Take a look at #485 and weigh in if you have comments! Will try to get that merged soon.

[mattjj]

@jekbradbury Hmm perhaps, but in this case couldn’t the masking stay at the level of the while loop primitive, rather than being pushed down into the body function (like a JAX interpreter would do)?