Implemented jax.lax.while primitive

[ymahlau]

This PR adds an implementation of the jax.lax.while primitive

Closes #15

[patrick-kidger]

Huh! This... actually seems really simple. I was not expecting it to be this easy.

I think to merge this I would like to see some tests added, if that's okay? The ones that jump out are:

• a basic test with quaxified carry.
• something where the cond_fun and body_fun close over additional constant values.
• combining this with jit/grad/vmap.

[ymahlau]

I agree that some test cases are necessary. To this end, I extended the unitful example and added it to the examples folder to have some wrapper that allows for easy testing. Then I added some test cases including constant values, jit and vmap.

[patrick-kidger]

Nice, thank you! My comments on this are:

1. it looks like the pre-commit hooks are failing (formatting etc.). These should be easy enough to fix: pre-commit install; pre-commit run --all-files; git add -u; git commit.
2. I have a very specific request: can you add some a test in which we (a) close over a variable in the body function whilst simultaneously (b) differentiating with respect to it? This is something which I have come to know can easily go wrong... so I'd like to be sure we get it right!