JAX (software)

<p>JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning. It is developed by <a href="/facts/Google/GT9Sugza">Google</a> with contributions from <a href="/facts/Nvidia/8QfCgqbu">Nvidia</a> and other community contributors.
</p><p>It is described as bringing together a modified version of <a href="https://github.com/HIPS/autograd">autograd</a> (automatic obtaining of the gradient function through differentiation of a function) and OpenXLA's <a href="/facts/Accelerated_Linear_Algebra/EqnFTdFp">XLA</a> (Accelerated Linear Algebra). It is designed to follow the structure and workflow of <a href="/facts/NumPy/T6FnhWWD">NumPy</a> as closely as possible and works with various existing frameworks such as <a href="/facts/TensorFlow/kDwfBeNd">TensorFlow</a> and <a href="/facts/PyTorch/4Gv1t13B">PyTorch</a>. The primary features of JAX are:
</p>
<ol><li>Providing a unified <a href="/facts/NumPy/T6FnhWWD">NumPy</a>-like interface to computations that run on CPU, GPU, or TPU, in local or distributed settings.</li>
<li>Built-in Just-In-Time (JIT) compilation via Open XLA, an open-source machine learning compiler ecosystem.</li>
<li>Efficient evaluation of gradients via its <a href="/facts/Automatic_differentiation/dQcdhBAM">automatic differentiation</a> transformations.</li>
<li><a href="/facts/Auto-vectorization/YzD5uHsX">Automatically vectorized</a> to efficiently map them over arrays representing batches of inputs.</li></ol>

JAX (software) open-in-new

JAX (software)