Fix jax_intro timeout: use lax.fori_loop instead of Python for loop

The compute_call_price_jax function was timing out during cache.yml builds
because JAX unrolls Python for loops during JIT compilation. With large
arrays (M=10M), this causes excessive compilation time.

Solution: Replace Python for loop with jax.lax.fori_loop, which compiles
the loop efficiently without unrolling.

Fixes cell execution timeout in jax_intro.md

Author: mmcky

lectures/jax_intro.md

@@ -832,16 +832,28 @@ def compute_call_price_jax(β=β,
 
     s = jnp.full(M, np.log(S0))
     h = jnp.full(M, h0)
-    for t in range(n):
+
+    def loop_body(i, state):
+        s, h, key = state
         key, subkey = jax.random.split(key)
         Z = jax.random.normal(subkey, (2, M))
         s = s + μ + jnp.exp(h) * Z[0, :]
         h = ρ * h + ν * Z[1, :]
+        return s, h, key
+
+    s, h, key = jax.lax.fori_loop(0, n, loop_body, (s, h, key))
+
     expectation = jnp.mean(jnp.maximum(jnp.exp(s) - K, 0))
 
     return β**n * expectation
 ```
 
+```{note}
+We use `jax.lax.fori_loop` instead of a Python `for` loop.
+This allows JAX to compile the loop efficiently without unrolling it,
+which significantly reduces compilation time for large arrays.
+```
+
 Let's run it once to compile it:
 
 ```{code-cell} ipython3