Wrestling with vmap in JAX

When I discovered JAX, it ticked all the right boxes for me: performant, flexible and extensible number crunching using a familiar API and a functional programming mindset. I was keen to move beyond the tutorial phase and test it with a real life application. I decided to port a personal project that I wrote almost 10 years ago in numpy and OpenCV to jax. The old project is called "ChopShop". It is a library for computational art.