runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=quantecon_ubuntu2404/disk=large"

- name: Check nvidia Drivers

shell: bash -l {0}

run: nvidia-smi

- name: Install JAX and Numpyro

shell: bash -l {0}

run: |

runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=quantecon_ubuntu2404/disk=large"

- name: Install JAX and Numpyro

shell: bash -l {0}

run: |

- name: Check nvidia drivers

shell: bash -l {0}

run: |

This lecture provides a short introduction to [Google JAX](https://github.com/jax-ml/jax).

JAX is a high-performance scientific computing library that provides

* a NumPy-like interface that can automatically parallize across CPUs and GPUs,

* a just-in-time compiler for accelerating a large range of numerical

operations, and

* automatic differentiation.

Increasingly, JAX also maintains and provides more specialized scientific

computing routines, such as those originally found in SciPy.

```{admonition} GPU

One of the attractive features of JAX is that, whenever possible, its array

processing operations conform to the NumPy API.

This means that, in many cases, we can use JAX is as a drop-in NumPy replacement.

On a GPU, this code runs much faster than its NumPy equivalent.

Also, typically, the second run is faster than the first due to JIT compilation.

The outcome is similar to the `cos` example --- JAX is faster, especially on the second run after JIT compilation.

```{admonition} GPU

Once compiled, JAX is significantly faster than NumPy due to GPU acceleration.

By avoiding the large input arrays `x_mesh` and `y_mesh`, this `vmap` version uses far less memory.

When run on a CPU, its runtime is similar to that of the meshgrid version.

When run on a GPU, it is usually significantly faster.

In fact, using `vmap` has another advantage: It allows us to break vectorization up into stages.

This leads to code that is often easier to comprehend than traditional vectorized code.

We will investigate these ideas more when we tackle larger problems.

While we avoid large input arrays in the preceding version,

Let's try a slightly different approach that takes the max to the inside.

Because of this change, we never compute the two-dimensional array `f(x,y)`.

Here

* `f_vec_x_max` computes the max along any given row

* `f_vec_max` is a vectorized version that can compute the max of all rows in parallel.

We apply this function to all rows and then take the max of the row maxes.

Let's try it.

If you are running this on a GPU, as we are, you should see another nontrivial speed gain.

This lecture series has access to the following GPU

```{code-cell} ipython

You can check the backend used by JAX using:

```{code-cell} ipython3

import jax

import jax.numpy as jnp

devices = jax.devices()

print(f"The available devices are: {devices}")

def matrix_multiply(a, b):

return jnp.dot(a, b)

key = jax.random.PRNGKey(0)

x = jax.random.normal(key, (1000, 1000))

y = jax.random.normal(key, (1000, 1000))

z = matrix_multiply(x, y)

print(z)

devices = jax.devices()

print(f"The available devices are: {devices}")