HARK jaxinterp

[alanlujan91]

I experimented on this a while ago but haven't done work on it since.

https://github.com/alanlujan91/HARK/blob/6fd7796402bac8527860d9bb57f151b29a5e059d/HARK/interpolation/jaxinterp.py

Also trying to overhaul HARK's interpolation infrastructure. A different route I went to use GPU was using cupy as seen in this PR econ-ark/HARK#1203

I am writing a flexible multilinear interpolator that can use different targets depending on user's hardware: cpu serial, cpu parallel, and gpu https://github.com/alanlujan91/HARK/blob/6fd7796402bac8527860d9bb57f151b29a5e059d/HARK/interpolation/jaxinterp.py

would be happy to collaborate

[alanlujan91]

some of my own benchmarks

[Smit-create]

Thanks @alanlujan91. Can you please show a small example to use your API so that we can include and benchmark your implementation using: https://gist.github.com/HumphreyYang/f524bfeb21f038d701f4d2b9a7fbb8e6?