search

QuantEcon / cbc_workshops

Workshops for the Central Bank of Chile
14 stars 13 forks source link

Improve Aiyagari JAX code #1

Open jstac opened 2 years ago

jstac commented 2 years ago

With fixed r and w (interest rate and wages), the household problem and computation of aggregate capital are fast in https://github.com/QuantEcon/cbc_workshops/blob/main/aiyagari/aiyagari_jax.ipynb.

(I'm computing on the GPU, using either my machine or colab.)

But when I try to search over different r and w for the equilibrium at the end, which means solving many household problems, JAX slows up.

Put differently, compute time increases exponentially in the number of times I solve the household problem, rather than linearly, as I would expect.

@Smit-create, I'm hoping we can work together to resolve this. I'm happy to explain over zoom.

This is an important benchmark macro model, so it would be good to get it working well.

CC @thomassargent30 and @cc7768, who like these kinds of problems.

Smit-create commented 2 years ago

@Smit-create, I'm hoping we can work together to resolve this. I'm happy to explain over zoom.

Thanks, we can have a look together.