Optimisations

[EtienneCmb]

• [x] vmap = On GPU, it works much better when it's vmapping over the last axis
• [ ] pmap = there's a way of doing multi-processing (check)
• [ ] use iterators instead of arrays. But does it work with Jax?