Through this post I’m gonna deep dive into one of the most exciting features of the JAX library: vmap, or vectorized map. For high-performance computing and machine learning research, efficiently handling computations over large datasets is critical. This is where JAX, a numerical computation library in Python developed by Google, comes in. Not only does […]