@@ -208,7 +208,7 @@ time `T`) corresponds to firm size distribution in (approximate) equilibrium.
208208
209209``` {code-cell} ipython3
210210def generate_cross_section(
211- firm, M=1_000_000 , T=500, s_init=1.0, seed=123
211+ firm, M=500_000 , T=500, s_init=1.0, seed=123
212212 ):
213213
214214 μ_a, σ_a, μ_b, σ_b, μ_e, σ_e, s_bar = firm
@@ -231,14 +231,16 @@ def generate_cross_section(
231231 return s
232232```
233233
234+ Let's try running the code and generating a cross-section.
235+
234236``` {code-cell} ipython3
235237firm = Firm()
236238tic()
237239data = generate_cross_section(firm).block_until_ready()
238240toc()
239241```
240242
241- Running the above function again so we can see the speed with and without compile time.
243+ We run the function again so we can see the speed without compile time.
242244
243245``` {code-cell} ipython3
244246tic()
@@ -264,13 +266,14 @@ The plot produces a straight line, consistent with a Pareto tail.
264266
265267#### Alternative implementation with ` lax.fori_loop `
266268
267- We did not JIT-compile the ` for ` loop above because
268- acceleration of outer loops makes relatively little difference terms of
269- compute time.
269+ Although we JIT-compiled some of the code above,
270+ we did not JIT-compile the ` for ` loop.
271+
272+ Let's try squeezing out a bit more speed
273+ by
270274
271- However, to maximize performance, let's try squeezing out a bit more speed
272- by replacing the ` for ` loop with
273- [ ` lax.fori_loop ` ] ( https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.fori_loop.html ) .
275+ * replacing the ` for ` loop with [ ` lax.fori_loop ` ] ( https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.fori_loop.html ) and
276+ * JIT-compiling the whole function.
274277
275278Here a the ` lax.fori_loop ` version:
276279
@@ -348,7 +351,7 @@ Try writing an alternative version of `generate_cross_section_lax()` where the e
348351
349352Does it improve the runtime?
350353
351- What are the pros and cons of this approach.
354+ What are the pros and cons of this approach?
352355
353356``` {exercise-end}
354357```
@@ -401,11 +404,14 @@ data = generate_cross_section_lax(firm).block_until_ready()
401404toc()
402405```
403406
404- This method might be faster in some cases but in general the
405- relative speed will depend on the size of the cross-section and the length of
407+ This method might or might not be faster.
408+
409+ In general, the relative speed will depend on the size of the cross-section and the length of
406410the simulation paths.
407411
408- Also, this method is far more memory intensive.
412+ However, this method is far more memory intensive.
413+
414+ It will fail when $T$ and $M$ become sufficiently large.
409415
410416``` {solution-end}
411- ```
417+ ```
0 commit comments