with qe.Timer(precision=8):

z_max_numpy = np.max(f(x, y))

print(f"NumPy result: {z_max_numpy}")

print(f"NumPy result: {z_max_numpy:.6f}")

```

In the vectorized version, all the looping takes place in compiled code.

Moreover, NumPy uses implicit multithreading, so that at least some parallelization occurs.

```{note}

If you have a system monitor such as htop (Linux/Mac) or perfmon

(Windows), then try running this and then observing the load on your CPUs.

(You will probably need to bump up the grid size to see large effects.)

The output typically shows that the operation is successfully distributed across multiple threads.

```

(The parallelization cannot be highly efficient because the binary is compiled

before it sees the size of the arrays `x` and `y`.)

grid = np.linspace(-3, 3, 3_000)

with qe.Timer(precision=8):

compute_max_numba(grid)

z_max_numpy = compute_max_numba(grid)

print(f"Numba result: {z_max_numpy:.6f}")

```

Let's run again to eliminate compile time.

```{code-cell} ipython3

with qe.Timer(precision=8):

compute_max_numba(grid)

```

Depending on your machine, the Numba version can be a bit slower or a bit faster

than NumPy.

```{code-cell} ipython3

z_max_parallel_incorrect = compute_max_numba_parallel(grid)

print(f"Incorrect parallel Numba result: {z_max_parallel_incorrect}")

print(f"NumPy result: {z_max_numpy}")

print(f"Numba result: {z_max_parallel_incorrect} 😱")

```

The incorrect parallel implementation typically returns `-inf` (the initial value of `m`) instead of the correct maximum value of approximately `0.9999979986680024`.

The reason is that the variable $m$ is shared across threads and not properly controlled.

The reason is that the variable `m` is shared across threads and not properly controlled.

When multiple threads try to read and write `m` simultaneously, they interfere with each other, causing a race condition.

When multiple threads try to read and write `m` simultaneously, they interfere with each other.

This results in lost updates—threads read stale values of `m` or overwrite each other's updates—and the variable often never gets updated from its initial value of `-inf`.

Threads read stale values of `m` or overwrite each other's updates --— or `m` never gets updated from its initial value.

Here's a more carefully written version.

Now the code block that `for i in numba.prange(n)` acts over is independent

across `i`.

Each thread writes to a separate element of the array `row_maxes`.

Hence the parallelization is safe.

Here's the timings.

Each thread writes to a separate element of the array `row_maxes` and

the parallelization is safe.

```{code-cell} ipython3

with qe.Timer(precision=8):

compute_max_numba_parallel(grid)

z_max_parallel = compute_max_numba_parallel(grid)

print(f"Numba result: {z_max_parallel:.6f}")

```

Here's the timing.

```{code-cell} ipython3

with qe.Timer(precision=8):

compute_max_numba_parallel(grid)

```

If you have multiple cores, you should see at least some benefits from parallelization here.

If you have multiple cores, you should see at least some benefits from

parallelization here.

For more powerful machines and larger grid sizes, parallelization can generate major speed gains, even on the CPU.

For more powerful machines and larger grid sizes, parallelization can generate

major speed gains, even on the CPU.

### Vectorized code with JAX

In most ways, vectorization is the same in JAX as it is in NumPy.

On the surface, vectorized code in JAX is similar to NumPy code.

But there are also some differences, which we highlight here.

x_mesh, y_mesh = np.meshgrid(grid, grid)

with qe.Timer(precision=8):

z_max = jnp.max(f(x_mesh, y_mesh)).block_until_ready()

z_max = jnp.max(f(x_mesh, y_mesh))

z_max.block_until_ready()

print(f"Plain vanilla JAX result: {z_max:.6f}")

```

Let's run again to eliminate compile time.

```{code-cell} ipython3

with qe.Timer(precision=8):

z_max = jnp.max(f(x_mesh, y_mesh)).block_until_ready()

z_max = jnp.max(f(x_mesh, y_mesh))

z_max.block_until_ready()

```

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

with qe.Timer(precision=8):

z_max = jnp.max(f_vec(grid))

z_max.block_until_ready()

print(f"JAX vmap v1 result: {z_max:.6f}")

```

```{code-cell} ipython3

```{code-cell} ipython3

with qe.Timer(precision=8):

z_max = compute_max_vmap_v2(grid).block_until_ready()

print(f"JAX vmap v1 result: {z_max:.6f}")

```

Let's run it again to eliminate compilation time:

In our view, JAX is the winner for vectorized operations.

It dominates NumPy both in terms of speed (via JIT-compilation and parallelization) and memory efficiency (via vmap).

It dominates NumPy both in terms of speed (via JIT-compilation and

parallelization) and memory efficiency (via vmap).

Moreover, the `vmap` approach can sometimes lead to significantly clearer code.

While Numba is impressive, the beauty of JAX is that, with fully vectorized

operations, we can run exactly the

same code on machines with hardware accelerators and reap all the benefits

without extra effort.

operations, we can run exactly the same code on machines with hardware

accelerators and reap all the benefits without extra effort.

Moreover, JAX already knows how to effectively parallelize many common array

operations, which is key to fast execution.

For almost all cases encountered in economics, econometrics, and finance, it is

For most cases encountered in economics, econometrics, and finance, it is

far better to hand over to the JAX compiler for efficient parallelization than to

try to hand code these routines ourselves.

```{note}

Sharp readers will notice that we specify `device=cpu` in the `jax.jit` decorator.

The computation consists of many very small `lax.scan` iterations that must run sequentially, leaving little opportunity for the GPU to exploit parallelism.

The computation consists of many small sequential operations, leaving little

opportunity for the GPU to exploit parallelism.

As a result, kernel-launch overhead tends to dominate on the GPU, making the CPU a better fit for this workload.

As a result, kernel-launch overhead tends to dominate on the GPU, making the CPU

a better fit for this workload.

Curious readers can try removing this option to see how performance changes.

```

x_jax = qm_jax(0.1, n).block_until_ready()

```

JAX is also efficient for this sequential operation.

JAX is also quite efficient for this sequential operation.

Both JAX and Numba deliver strong performance after compilation, with Numba

typically (but not always) offering slightly better speeds on purely sequential

operations.

### Summary

While both Numba and JAX deliver strong performance for sequential operations,

there are significant differences in code readability and ease of use.

*there are significant differences in code readability and ease of use*.

The Numba version is straightforward and natural to read: we simply allocate an

array and fill it element by element using a standard Python loop.

For this type of sequential operation, Numba is the clear winner in terms of

code clarity and ease of implementation, as well as high performance.