> JAX is a Python library which augments numpy and Python code with function transformations which make it trivial to perform operations common in machine learning programs.