Jit the returned lcm function

OpenSourceEconomics / lcm

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

Apache License 2.0

16 stars 1 forks source link

Jit the returned lcm function #77

Closed mj023 closed 5 months ago

mj023 commented 5 months ago

In this PR we want to jit the function that is returned when we tell lcm to solve a model. This speeds up the execution time when solving a model for the first time and greatly reduces the execution time when solving the same model with different parameters.

[X] Let Jax use the gpu backend if available
[X] Jit the returned lcm function