Precompile CCV Functions

OpenSourceEconomics / lcm

Solution and simulation of life cycle models in Python with GPU acceleration.

Apache License 2.0

16 stars 1 forks source link

Precompile CCV Functions #73

Closed mj023 closed 5 months ago

mj023 commented 6 months ago

[X] Let Jax use the gpu backend if available
[x] Remove inefficient Jit Operations
[x] Concurrently precompile all functions so that they are ready when needed in each period

mj023 commented 5 months ago

It turns out, that just jitting the function that get_lcm_function() returns increases the speed for one iteration, while simultaneously drastically increasing the speed if you want to solve the model multiple times with different parameters e.g. the parameter estimation use case. While precompiling is still faster for the first iteration, it is not possible to both precompile the functions inside the returned function and then jit everything. Therefore I will close this pull request.

timmens commented 5 months ago

Regarding the previous comment, see #77