Add compilation cache support

Caltech-AMBER / ambersim

In-house tools built on GPU-accelerated simulation

MIT License

7 stars 2 forks source link

Add compilation cache support #8

Open alberthli opened 10 months ago

alberthli commented 10 months ago

It would be extremely useful to add utils + examples for using jax's compilation_cache feature, which could really speed up experiments over time when using the same machine.