vmap support / batching rule for lax.cond

[jonasrauber]

It would be great to have vmap support for lax.cond similar to the recently added vmap support for lax.while_loop. I know that I can use np.where as a workaround for some cases, but for computationally expensive if conditions a batched cond might be better.

[jonasrauber]

As I usually only need the if condition (no else part) I am considering a single iteration while_loop instead of cond as another workaround. What do you think?

[mattjj]

I think this feature sounds super fun to implement and I can't wait to try it!

But in the meantime a single-iteration while_loop (or `fori_loop if it's more convenient) is probably a good workaround.