Add JAX

[amalrkc]

JAX

Google JAX is a machine learning framework for transforming numerical functions, to be used in Python. It is described as bringing together a modified version of autograd (automatic obtaining of the gradient function through differentiation of a function) and TensorFlow's XLA (Accelerated Linear Algebra). It is designed to follow the structure and workflow of NumPy as closely as possible and works with various existing frameworks such as TensorFlow and PyTorch. The primary functions of JAX are:

1. grad: automatic differentiation
2. jit: compilation
3. vmap: auto-vectorization
4. pmap: SPMD programming

Wikipedia

Github Repo

JAX in 100 seconds by Fireship

Link to the vector file