IAML Blog

JAX, aka NumPy on steroids

14, Jan Simone Scardapane / trad. Stefano Di Pietro

In the age of the 'big ones' (TensorFlow, PyTorch, ...), introducing and studying a new machine learning library might seem counterproductive. Yet JAX, a brand new research project by Google, has several features that make it interesting to a large audience. Firstly, it looks like a NumPy wrapper, making the transition from this library almost immediate. Secondly, it makes efficiency one of its strengths, thanks to the transparent use of XLA, a linear algebra accelerator originall…

Continue Reading...