FIX: jax_intro timeout: use lax.fori_loop instead of Python for loop (#442)

* 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

* style: use jstac's fori_loop naming conventions

- loop_body -> update
- state -> loop_state
- Added explicit new_loop_state and final_loop_state variables
- More verbose but clearer for first-time fori_loop readers

* style: loop_state -> initial_loop_state

Author: mmcky

lectures/jax_intro.md

@@ -832,16 +832,31 @@ def compute_call_price_jax(β=β,
 
     s = jnp.full(M, np.log(S0))
     h = jnp.full(M, h0)
-    for t in range(n):
+
+    def update(i, loop_state):
+        s, h, key = loop_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, :]
+        new_loop_state = s, h, key
+        return new_loop_state
+
+    initial_loop_state = s, h, key
+    final_loop_state = jax.lax.fori_loop(0, n, update, initial_loop_state)
+    s, h, key = final_loop_state
+
     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