TST: test timing of jax

QuantEcon · Mar 22, 2024 · 437c7db · 437c7db
1 parent 1a1be32
commit 437c7db
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/lectures/jax_intro.md b/lectures/jax_intro.md
@@ -375,7 +375,11 @@ f_jit(x)
 And now let's time it.
 
 ```{code-cell} ipython3
-%time f_jit(x).block_until_ready()
+%time x_jax = jax.device_put(x)
+```
+
+```{code-cell} ipython3
+%timeit f_jit(x_jax).block_until_ready()
 ```
 
 Note the speed gain.