JAX: Add JAX implementation for wealth dynamics by Smit-create · Pull Request #304 · QuantEcon/lecture-python.myst

Smit-create · 2022-12-30T10:51:44Z

github-actions · 2022-12-30T11:11:42Z

🚀 Deployed on https://63c622d321fc755cc8652567--nostalgic-wright-5fa355.netlify.app

Smit-create · 2022-12-31T06:15:30Z

This seems to be ready for the first review. Please have a look. Thanks

HumphreyYang

Many Thanks for the update @Smit-create. The code looks great to me. One improvement we could potentially make is to hide the code using numba using

---
tags: [hide-input]
---

and only compare the speed when making the final comparison.

I think it is clear to readers when comparing the performance of two implementations in my case since there is only one function involved, but in this case, having two implementations throughout the lecture may confuse readers and obscure the focus of the lecture. Please let me know what you think :) (also CC @jstac)

HumphreyYang · 2023-01-01T09:18:24Z

+# Check if JAX is using GPU
+print("jax backend: {}".format(jax.lib.xla_bridge.get_backend().platform))
+# Check the devices available for JAX
+print(jax.devices())


I think we can use

print(f"jax backend: {jax.devices()[0].platform}")

in our PRs for simplicity and consistency.

Thanks @HumphreyYang , simple is good.

@Smit-create , please change jax to JAX in the text.

Thanks, I will change that accordingly.

Smit-create · 2023-01-04T06:48:52Z

Please let me know what you think :)

This sounds like a nice idea, let's try that out and see how it looks.

jstac · 2023-01-09T00:26:43Z

@Smit-create and @HumphreyYang, please let me know when this is ready for another review.

HumphreyYang · 2023-01-09T01:24:28Z

Please let me know what you think :)

This sounds like a nice idea, let's try that out and see how it looks.

@Smit-create I think it looks cleaner and simpler with Numba code hidden. How do you feel about this change?

Smit-create · 2023-01-09T06:03:10Z

How do you feel about this change?

I think that looks cleaner. What do you think @jstac?

please let me know when this is ready for another review.

@jstac, I think that the content is ready for review. Might just need to edit/add a couple of cells following the discussions in #301 (comment).

jstac · 2023-01-10T05:12:59Z

-## Implementation
+## Implementation using JAX
+
+Let's define a Model to represent the wealth dynamics.


Please change "Model" to "model" and the same below.

jstac · 2023-01-10T05:13:58Z

+update_states_jax = jax.jit(update_states_jax, static_argnums=(2,))
+```
+
+Here’s function to simulate the time series of wealth for individual households using `for` loop and JAX.


using a for loop

jstac · 2023-01-10T05:15:23Z

+
+# 2nd time is expected to be very fast because of JIT
+w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size)
+```


does this need block_until_ready?

I guess not but just checking.

Let's add that. The documentation for benchmarking also recommends using it: https://jax.readthedocs.io/en/latest/faq.html#benchmarking-jax-code

jstac · 2023-01-10T05:18:42Z

Looking at the current deployment https://63b52771da28e92bd5e7733b--nostalgic-wright-5fa355.netlify.app/wealth_dynamics.html shows that one numba code cell is not hidden, while the others are. Is this intensional?

jstac · 2023-01-10T05:20:52Z

Should we compare the timing of update_cross_section_jax and update_cross_section?

I suspect that, on a large cross-section, the JAX code will be much faster...

jstac · 2023-01-10T05:21:36Z

Does this sentence need to be adjusted/removed?

"Notice how effectively Numba has implemented multithreading for this routine: all 8 CPUs on our workstation are running at maximum capacity (even though four of them are virtual)."

I think we are using JAX at this point.

jstac · 2023-01-10T05:23:16Z

Many thanks @Smit-create . Please see my comments above.

@mmcky , if you have time please have a quick glance at this once @Smit-create has pushed changes.

Smit-create · 2023-01-10T05:57:23Z

shows that one numba code cell is not hidden, while the others are. Is this intentional?

I thought to keep it unhidden so that the below cell that measures time can relate to it. But let's hide all the numba code and just keep the JAX available readily.

Should we compare the timing of update_cross_section_jax and update_cross_section?

It's already been compared using generate_lorenz_and_gini_jax (while plotting) and shows about a ~30s difference.

"Notice how effectively Numba has implemented multithreading for this routine: all 8 CPUs on our workstation are running at maximum capacity (even though four of them are virtual)."

Let's remove this completely along with the image?

mmcky · 2023-01-13T02:10:20Z

thanks @Smit-create, @HumphreyYang and @jstac

Is the purpose of hiding the Numba code to clean up the lecture and enable those interested to see the code. It is not the approach often taken relative to our other lectures which tends to always show the code but I see the application here.

This is what it looks like in the preview

mmcky · 2023-01-13T02:12:08Z

For the text in the lecture

Notice the time difference between the wealth_time_series and wealth_time_series_jax

Could we use the glu feature to have the values here? It is a statement that requires the reader to find the values and compare them. https://jupyterbook.org/en/stable/reference/cheatsheet.html?highlight=glu#gluing-numbers

jstac · 2023-01-13T02:51:53Z

Yes, I think it looks nice to have the numba code hidden, and I can understand the motivation, but at the same time it might be confusing for the reader --- why is some code hidden and some not?

Also, if we hide numba code here, then we should do it elsewhere when we use JAX --- but then it's a burden and possibly unclear exactly what we should hide.

Having thought all this through, I suggest that we show the numba code rather than hiding it and do the same for the remaining lectures.

Smit-create · 2023-01-13T15:19:18Z

I suggest that we show the numba code rather than hiding it and do the same for the remaining lectures.

Sure, will do that.

Could we use the glue feature to have the values here?

Yes, added it, thanks!

jstac · 2023-01-13T22:05:35Z

Thanks @Smit-create

@HumphreyYang just checking that you saw the discussion above regarding hiding / not hiding.

Could you please give this PR one more careful look over and then let me know when you are done. After that I think we'll be ready to merge.

HumphreyYang · 2023-01-13T23:19:07Z

Thanks @Smit-create

@HumphreyYang just checking that you saw the discussion above regarding hiding / not hiding.

Could you please give this PR one more careful look over and then let me know when you are done. After that I think we'll be ready to merge.

Hi @jstac,

Before that, could you please have a look at the GPU Warning and JAX imports I pushed in my PR so that @Smit-create and I can have a consistent GPU warning?

jstac · 2023-01-14T00:21:37Z

Thanks @HumphreyYang , good idea. Here's what I suggest for the GPU warning, which slightly modifies the version of @HumphreyYang :

This lecture is built using hardware that has access to a GPU and uses JAX for GPU programming. As a result, the lecture will be slower when running on a machine without a GPU.

Free GPUs are available on Google Colab. To use this option, please click on the play icon top right, select Colab, and set the runtime environment to include a GPU.

Alternatively, if if you have your own GPU, you can follow the instructions for for installing JAX with GPU support.

Smit-create · 2023-01-16T06:15:56Z

The PR is ready now. The failure is because of this diff. Once we merge that PR, I'll rebase this on top of that and that should fix the CI.

jstac · 2023-01-16T17:20:50Z

Thanks @Smit-create

Please see the requested changes to GPU Warning in [kesten_processes] Rewrite Main Exercise Using JAX #301 --- sorry for repeated changes and thanks for your patience
Is there a working netlify deployment I can look at here?

Smit-create · 2023-01-17T04:03:43Z

Please see the requested changes to GPU Warning

Thanks for the review @jstac. I have updated the diff.

Is there a working netlify deployment I can look at here?

The last commit (10th commit) was failing on CI and so I suppose it won't be deployed. I have pushed a new commit and let's see if that works.

Smit-create · 2023-01-17T07:52:08Z

The tests pass and the deployed version is available at https://63c622d321fc755cc8652567--nostalgic-wright-5fa355.netlify.app/

jstac · 2023-01-17T18:39:35Z

Great work @Smit-create , many thanks. Merging.

add jax implementation for wealth dynamics time series

e6699ee

Add update cross section in jax

9e1657e

Smit-create changed the title ~~JAX: Add JAX implementation for wealth dynamics #303~~ JAX: Add JAX implementation for wealth dynamics Dec 31, 2022

Smit-create added 2 commits December 31, 2022 10:48

fix an error

a23e692

add a cell to check backend

2cb8379

Smit-create requested a review from HumphreyYang December 31, 2022 06:14

Smit-create marked this pull request as ready for review December 31, 2022 06:15

HumphreyYang reviewed Jan 1, 2023

View reviewed changes

Smit-create added 2 commits January 4, 2023 09:55

Apply suggestions

ff9d4c6

Add for-loop implementation

063367f

Hide numba code

d942383

jstac reviewed Jan 10, 2023

View reviewed changes

Hide all numba code and fix typos

e43c237

Remove hidden tags and use glue

975be52

Add GPU warning

6b49755

Update GPU warning

d4e56ba

jstac merged commit 5f7ea92 into main Jan 17, 2023

jstac deleted the jax_wealth_dynamics_1 branch January 17, 2023 18:39

Smit-create mentioned this pull request Apr 18, 2023

Remove JAX from Wealth dynamics #335

Merged

Uh oh!

Uh oh!

Conversation

Smit-create commented Dec 30, 2022

Uh oh!

github-actions Bot commented Dec 30, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Smit-create commented Dec 31, 2022

Uh oh!

HumphreyYang left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HumphreyYang Jan 1, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jstac Jan 2, 2023

Choose a reason for hiding this comment

Uh oh!

Smit-create Jan 3, 2023

Choose a reason for hiding this comment

Uh oh!

Smit-create commented Jan 4, 2023

Uh oh!

jstac commented Jan 9, 2023

Uh oh!

HumphreyYang commented Jan 9, 2023

Uh oh!

Smit-create commented Jan 9, 2023

Uh oh!

jstac Jan 10, 2023

Choose a reason for hiding this comment

Uh oh!

jstac Jan 10, 2023

Choose a reason for hiding this comment

Uh oh!

jstac Jan 10, 2023

Choose a reason for hiding this comment

Uh oh!

Smit-create Jan 10, 2023

Choose a reason for hiding this comment

Uh oh!

jstac commented Jan 10, 2023

Uh oh!

jstac commented Jan 10, 2023

Uh oh!

jstac commented Jan 10, 2023

Uh oh!

jstac commented Jan 10, 2023

Uh oh!

Smit-create commented Jan 10, 2023

Uh oh!

mmcky commented Jan 13, 2023

Uh oh!

mmcky commented Jan 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jstac commented Jan 13, 2023

Uh oh!

Smit-create commented Jan 13, 2023

Uh oh!

jstac commented Jan 13, 2023

Uh oh!

HumphreyYang commented Jan 13, 2023

Uh oh!

jstac commented Jan 14, 2023

Uh oh!

Smit-create commented Jan 16, 2023

Uh oh!

jstac commented Jan 16, 2023

Uh oh!

Smit-create commented Jan 17, 2023

Uh oh!

Smit-create commented Jan 17, 2023

Uh oh!

jstac commented Jan 17, 2023

Uh oh!

Reviewers

github-actions Bot commented Dec 30, 2022 •

edited

Loading

HumphreyYang left a comment •

edited

Loading

HumphreyYang Jan 1, 2023 •

edited

Loading

mmcky commented Jan 13, 2023 •

edited

Loading