update %timeit

Author: HumphreyYang

lectures/jax_intro.md

@@ -291,23 +291,47 @@ x = jnp.ones(n)
 How long does the function take to execute?
 
 ```{code-cell} ipython3
-%time f(x).block_until_ready()
+%timeit f(x).block_until_ready()
 ```
 
 ```{note}
 Here, in order to measure actual speed, we use the `block_until_ready()` method 
 to hold the interpreter until the results of the computation are returned from
 the device. This is necessary because JAX uses asynchronous dispatch, which
 allows the Python interpreter to run ahead of GPU computations.
+```
+
+```{note}
+Here, we use the [`%timeit` magic](https://ipython.readthedocs.io/en/stable/interactive/magics.html#magic-timeit)
+to time the execution of the function. 
+
+This command runs the code multiple times to get a more accurate measurement.
+
+Alternatively, we could use the `%time` magic command:
+
+```{code-cell} ipython3
+%time f(x).block_until_ready()
+```
+
+Unlike `%timeit`, this command runs the code only once.
+
+The `%timeit` magic command offers several advantages over `%time`:
 
+* It executes the code multiple times, providing a more accurate measurement
+  for measuring short code snippets
+* It reports both the average execution time and standard deviation
+
+However, `%timeit` can be time-consuming for large computations.
+
+This is why we switch to using `%time` in later lectures.
 ```
 
 The code doesn't run as fast as we might hope, given that it's running on a GPU.
 
 But if we run it a second time it becomes much faster:
 
 ```{code-cell} ipython3
-%time f(x).block_until_ready()
+%timeit f(x).block_until_ready()
 ```
 
 This is because the built in functions like `jnp.cos` are JIT compiled and the
@@ -330,7 +354,7 @@ y = jnp.ones(m)
 ```
 
 ```{code-cell} ipython3
-%time f(y).block_until_ready()
+%timeit f(y).block_until_ready()
 ```
 
 Notice that the execution time increases, because now new versions of 
@@ -341,14 +365,14 @@ If we run again, the code is dispatched to the correct compiled version and we
 get faster execution.
 
 ```{code-cell} ipython3
-%time f(y).block_until_ready()
+%timeit f(y).block_until_ready()
 ```
 
 The compiled versions for the previous array size are still available in memory
 too, and the following call is dispatched to the correct compiled code.
 
 ```{code-cell} ipython3
-%time f(x).block_until_ready()
+%timeit f(x).block_until_ready()
 ```
 
 ### Compiling the outer function
@@ -368,7 +392,7 @@ f_jit(x)
 And now let's time it.
 
 ```{code-cell} ipython3
-%time f_jit(x).block_until_ready()
+%timeit f_jit(x).block_until_ready()
 ```
 
 Note the speed gain.
@@ -523,7 +547,7 @@ z_loops = np.empty((n, n))
 ```
 
 ```{code-cell} ipython3
-%%time
+%%timeit
 for i in range(n):
     for j in range(n):
         z_loops[i, j] = f(x[i], y[j])
@@ -564,14 +588,14 @@ x_mesh, y_mesh = jnp.meshgrid(x, y)
 Now we get what we want and the execution time is very fast.
 
 ```{code-cell} ipython3
-%%time
+%%timeit
 z_mesh = f(x_mesh, y_mesh).block_until_ready()
 ```
 
 Let's run again to eliminate compile time.
 
 ```{code-cell} ipython3
-%%time
+%%timeit
 z_mesh = f(x_mesh, y_mesh).block_until_ready()
 ```
 
@@ -591,14 +615,14 @@ x_mesh, y_mesh = jnp.meshgrid(x, y)
 ```
 
 ```{code-cell} ipython3
-%%time
+%%timeit
 z_mesh = f(x_mesh, y_mesh).block_until_ready()
 ```
 
 Let's run again to get rid of compile time.
 
 ```{code-cell} ipython3
-%%time
+%%timeit
 z_mesh = f(x_mesh, y_mesh).block_until_ready()
 ```
 
@@ -637,14 +661,14 @@ f_vec = jax.vmap(f_vec_y, in_axes=(0, None))
 With this construction, we can now call the function $f$ on flat (low memory) arrays.
 
 ```{code-cell} ipython3
-%%time
+%%timeit
 z_vmap = f_vec(x, y).block_until_ready()
 ```
 
 We run it again to eliminate compile time.
 
 ```{code-cell} ipython3
-%%time
+%%timeit
 z_vmap = f_vec(x, y).block_until_ready()
 ```
 
@@ -716,14 +740,14 @@ def compute_call_price_jax(β=β,
 Let's run it once to compile it:
 
 ```{code-cell} ipython3
-%%time 
+%%timeit 
 compute_call_price_jax().block_until_ready()
 ```
 
 And now let's time it:
 
 ```{code-cell} ipython3
-%%time 
+%%timeit 
 compute_call_price_jax().block_until_ready()
 ```