Uploaded February 2023 | Updated September 2026, 2 weeks ago
In the last two episodes of JAX Foundations, we gave you insight on what JAX is and who uses it. In this episode, we dive deeper, giving you insight on how you can use it. We’ll specifically show you how to evaluate JAX with Google Colab, seeing how it demonstrates derivatives of functions with jax.grad, Auto vectorization with jax.vmap, just-in-time compilation, and more. Watch this video to see how you can get started with JAX!
Resources:
JAX Quickstart Colab → https://goo.gle/3I8eVVe
Chapters:
0:00 - Intro
0:25 - How you can use JAX in Colab?
0:44 - How does JAX optimize mathematical operations?
1:38 - How can you download JAX onto your own work station?
1:50 - How to install JAX for CPU’s
2:14 -How to install Jax for GPU, CUDA, cuDNN, and Google Cloud TPU’s
2:47 - Leave us questions and comments below!
Watch more episodes of JAX Foundations → https://goo.gle/JAXFoundations
Subscribe to the Google Research Channel → https://goo.gle/GoogleResearch
In the last two episodes of JAX Foundations, we gave you insight on what JAX is and who uses it. In this episode, we dive deeper, giving you insight on how you can use it. We’ll specifically show you how to evaluate JAX with Google Colab, seeing how it demonstrates derivatives of functions with jax.grad, Auto vectorization with jax.vmap, just-in-time compilation, and more. Watch this video to see how you can get started with JAX!
Resources:
JAX Quickstart Colab → https://goo.gle/3I8eVVe
Chapters:
0:00 - Intro
0:25 - How you can use JAX in Colab?
0:44 - How does JAX optimize mathematical operations?
1:38 - How can you download JAX onto your own work station?
1:50 - How to install JAX for CPU’s
2:14 -How to install Jax for GPU, CUDA, cuDNN, and Google Cloud TPU’s
2:47 - Leave us questions and comments below!
Watch more episodes of JAX Foundations → https://goo.gle/JAXFoundations
Subscribe to the Google Research Channel → https://goo.gle/GoogleResearch
