[mccall_model_with_separation] Conversion to JAX - #501
Conversation
|
Hi John, @jstac , I just compare the numba version and the jax version. I find that Jax version is correct, and numba version has a small problem. In the original numba code, we didn't put everything in functions, so we have in the first loop over unemployment compensation c from 2 to 12. So the model in the numba version when looping over discount factor and separation rate, we have c=12 instead of the default value of 6. Similar happens when we loop discount factor openly without the function, but since the last loop is 0.98 over 0.99 the difference is not too big. So the conclusion is that the Jax version is correct and better. Best, |
longye-tian
left a comment
longye-tian
left a comment
There was a problem hiding this comment.
Jax version is correct
| w_bar = compute_reservation_wage(model) | ||
| return w_bar | ||
|
|
||
| for i, c in enumerate(c_vals): |
There was a problem hiding this comment.
we didn't put this part in the function.
So later one when we loop over discount factor, we have c=12 which is the value from the last loop
| w_bar = compute_reservation_wage(model) | ||
| return w_bar | ||
|
|
||
| for i, β in enumerate(beta_vals): |
| w_bar = compute_reservation_wage(model) | ||
| return w_bar | ||
|
|
||
| for i, α in enumerate(alpha_vals): |
|
Many thanks @jstac and @longye-tian, I agree. One way to verify that this is the issue is to add I got the same solution as JAX once we do that: |
|
@longye-tian and @HumphreyYang outstanding work.
|
|
Thanks @mmcky, Would you like me to open a PR that fix the numba code? Just curious if we are replacing the numba code with JAX code. |
Ah. Thank you @HumphreyYang. I am rushing around the repos at the moment. You're 100% correct. This is replacing the |
|
Thanks @HumphreyYang and @longye-tian for your helpful reviews. @jstac just checking do you intend to make any other changes before merge? |
|
Many thanks @longye-tian for the careful review, it's much appreciated. Thanks also @HumphreyYang for confirming and @mmcky for supervising. As @HumphreyYang says, we only need the JAX version. Can someone please do one more overall sanity check and make sure everything makes sense. Then please flag it for merging. I don't need to see it again. |
HumphreyYang
changed the title
4 participants
Changing job search with separation to use JAX
@longye-tian @HumphreyYang
Could one of you please review this? The last figure looks a bit different.