Where to learn how vmap works under the hood (vectorization in general?)

[PaulScemama]

I'd like to learn about how vmap actually works under the hood. I'm sure there is a place to learn about how it works and how vectorization in general works to speed up computation, but I haven't found it yet.

The level of granularity I'd like to know about vmap is to the extent that I can then implement my own (narrow and simple) vmap from scratch.

Any ideas? @jakevdp

[jakevdp]

For this, the best place to start would be autodidax, which is a stripped-down version of JAX's transformation system built from scratch. In particular, the autodidax implementation of vmap is here: https://jax.readthedocs.io/en/latest/autodidax.html#vectorized-batching-with-vmap

[PaulScemama]

Thanks!