Add @jax.jit decorators to VFI and OPI functions

- Improves performance by JIT compiling the main iteration routines
- Speedup now consistently 2.5-2.6x (vs 2.4-2.5x before)
- Absolute times also improved (OPI: ~0.3-0.4s vs ~0.4-0.5s)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: jstac

lectures/ifp_opi.md

@@ -222,6 +222,7 @@ def iterate_policy_operator(σ, v, m, model):
 For comparison, here's VFI from {doc}`ifp_discrete`:
 
 ```{code-cell} ipython3
+@jax.jit
 def value_function_iteration(model, tol=1e-5, max_iter=10_000):
     """
     Implements VFI using successive approximation.
@@ -252,6 +253,7 @@ The algorithm alternates between
 2. Computing a new greedy policy based on the updated value function
 
 ```{code-cell} ipython3
+@jax.jit
 def optimistic_policy_iteration(model, m=10, tol=1e-5, max_iter=10_000):
     """
     Implements optimistic policy iteration with step size m.