JAX (software)

<h2 id="grad">grad</h2>
<p class="note">Main article: <a href="/facts/Automatic_differentiation/dQcdhBAM">Automatic differentiation</a></p>
<p>The below code demonstrates the grad function's automatic differentiation. 
</p>
# imports
from jax import grad
import jax.numpy as jnp

# define the logistic function
def logistic(x):  
    return jnp.exp(x) / (jnp.exp(x) + 1)

# obtain the gradient function of the logistic function
grad_logistic = grad(logistic)

# evaluate the gradient of the logistic function at x = 1 
grad_log_out = grad_logistic(1.0)   
print(grad_log_out)

<p>The final line should outputː
</p>
0.19661194

<h2 id="jit">jit</h2>
<p>The below code demonstrates the jit function's optimization through fusion.
</p>
# imports
from jax import jit
import jax.numpy as jnp

# define the cube function
def cube(x):
    return x * x * x

# generate data
x = jnp.ones((10000, 10000))

# create the jit version of the cube function
jit_cube = jit(cube)

# apply the cube and jit_cube functions to the same data for speed comparison
cube(x)
jit_cube(x)

<p>The computation time for jit_cube (line #17) should be noticeably shorter than that for cube (line #16). Increasing the values on line #7, will further exacerbate the difference.
</p>
<h2 id="vmap">vmap</h2>
<p class="note">Main article: <a href="/facts/Array_programming/gFx1SxIY">Array programming</a></p>
<p>The below code demonstrates the vmap function's vectorization.
</p>
# imports
from jax import vmap partial
import jax.numpy as jnp

# define function
def grads(self, inputs):
    in_grad_partial = jax.partial(self._net_grads, self._net_params)
    grad_vmap = jax.vmap(in_grad_partial)
    rich_grads = grad_vmap(inputs)
    flat_grads = np.asarray(self._flatten_batch(rich_grads))
    assert flat_grads.ndim == 2 and flat_grads.shape[0] == inputs.shape[0]
    return flat_grads

<p>The GIF on the right of this section illustrates the notion of vectorized addition.
</p>

<h2 id="pmap">pmap</h2>
<p>The below code demonstrates the pmap function's parallelization for matrix multiplication.
</p>
# import pmap and random from JAX; import JAX NumPy
from jax import pmap, random
import jax.numpy as jnp

# generate 2 random matrices of dimensions 5000 x 6000, one per device
random_keys = random.split(random.PRNGKey(0), 2)
matrices = pmap(lambda key: random.normal(key, (5000, 6000)))(random_keys)

# without data transfer, in parallel, perform a local matrix multiplication on each CPU/GPU
outputs = pmap(lambda x: jnp.dot(x, x.T))(matrices)

# without data transfer, in parallel, obtain the mean for both matrices on each CPU/GPU separately
means = pmap(jnp.mean)(outputs)
print(means)

<p>The final line should print the valuesː
</p>
[1.1566595 1.1805978]

<h2 id="see-also">See also</h2>
<ul><li><a href="/facts/NumPy/T6FnhWWD">NumPy</a></li>
<li><a href="/facts/TensorFlow/kDwfBeNd">TensorFlow</a></li>
<li><a href="/facts/PyTorch/4Gv1t13B">PyTorch</a></li>
<li><a href="/facts/CUDA/fvPpolot">CUDA</a></li>
<li><a href="/facts/Accelerated_Linear_Algebra/EqnFTdFp">Accelerated Linear Algebra</a></li></ul>
<h2 id="external-links">External links</h2>
<ul><li>Documentationː <a href="https://jax.readthedocs.io/">jax.readthedocs.io</a></li>
<li>Colab (<a href="/facts/Project_Jupyter/sN94yfV5">Jupyter</a>/iPython) Quickstart Guideː <a href="https://colab.research.google.com/github/google/jax/blob/main/docs/notebooks/quickstart.ipynb">colab.research.google.com/github/google/jax/blob/main/docs/notebooks/quickstart.ipynb</a></li>
<li><a href="/facts/TensorFlow/kDwfBeNd">TensorFlow</a>'s XLAː <a href="https://www.tensorflow.org/xla">www.tensorflow.org/xla</a> (Accelerated Linear Algebra)</li>
<li><a href="/facts/YouTube/db276jst">YouTube</a> TensorFlow Channel "Intro to JAX: Accelerating Machine Learning research": <a href="https://www.youtube.com/watch?v=WdTeDXsOSj4">www.youtube.com/watch?v=WdTeDXsOSj4</a></li>
<li>Original paperː <a href="https://mlsys.org/Conferences/doc/2018/146.pdf">mlsys.org/Conferences/doc/2018/146.pdf</a></li></ul>

<h2 id="references">References</h2>

<ol>
<li id="fn:1"><p>Bradbury, James; Frostig, Roy; Hawkins, Peter; Johnson, Matthew James; Leary, Chris; MacLaurin, Dougal; Necula, George; Paszke, Adam; Vanderplas, Jake; Wanderman-Milne, Skye; Zhang, Qiao (2022-06-18), "JAX: Autograd and XLA", Astrophysics Source Code Library, Google, Bibcode:2021ascl.soft11002B, archived from the original on 2022-06-18, retrieved 2022-06-18 <a href="https://web.archive.org/web/20220618205214/https://github.com/google/jax" target="_blank">https://web.archive.org/web/20220618205214/https://github.com/google/jax</a> <a href="#fnref:1" class="footnote-back-ref">↩</a></p></li>
<li id="fn:2"><p>Frostig, Roy; Johnson, Matthew James; Leary, Chris (2018-02-02). "Compiling machine learning programs via high-level tracing" (PDF). MLsys: 1–3. Archived (PDF) from the original on 2022-06-21.{{cite journal}}:  CS1 maint: date and year (link) <a href="https://mlsys.org/Conferences/doc/2018/146.pdf" target="_blank">https://mlsys.org/Conferences/doc/2018/146.pdf</a> <a href="#fnref:2" class="footnote-back-ref">↩</a></p></li>
<li id="fn:3"><p>"Using JAX to accelerate our research". www.deepmind.com. Archived from the original on 2022-06-18. Retrieved 2022-06-18. <a href="https://www.deepmind.com/blog/using-jax-to-accelerate-our-research" target="_blank">https://www.deepmind.com/blog/using-jax-to-accelerate-our-research</a> <a href="#fnref:3" class="footnote-back-ref">↩</a></p></li>
<li id="fn:4"><p>Lynley, Matthew. "Google is quietly replacing the backbone of its AI product strategy after its last big push for dominance got overshadowed by Meta". Business Insider. Archived from the original on 2022-06-21. Retrieved 2022-06-21. <a href="https://web.archive.org/web/20220621143905/https://www.businessinsider.com/facebook-pytorch-beat-google-tensorflow-jax-meta-ai-2022-6" target="_blank">https://web.archive.org/web/20220621143905/https://www.businessinsider.com/facebook-pytorch-beat-google-tensorflow-jax-meta-ai-2022-6</a> <a href="#fnref:4" class="footnote-back-ref">↩</a></p></li>
<li id="fn:5"><p>"Why is Google's JAX so popular?". Analytics India Magazine. 2022-04-25. Archived from the original on 2022-06-18. Retrieved 2022-06-18. <a href="https://analyticsindiamag.com/why-is-googles-jax-so-popular/" target="_blank">https://analyticsindiamag.com/why-is-googles-jax-so-popular/</a> <a href="#fnref:5" class="footnote-back-ref">↩</a></p></li>
<li id="fn:6"><p>"Quickstart — JAX documentation". <a href="https://docs.jax.dev/en/latest/quickstart.html" target="_blank">https://docs.jax.dev/en/latest/quickstart.html</a> <a href="#fnref:6" class="footnote-back-ref">↩</a></p></li>
</ol>

JAX (software) open-in-new

JAX (software)