Closed jeremiecoullon closed 3 years ago
Since I put an emphasis on performance I should have benchmarks that are run with the CI at some point. Something like: https://github.com/marketplace/actions/continuous-benchmark
I did some quick performance tests (on the example in the reamde):
• the "non-hashable type" and the kernel_factory
fixes made no difference in performance (running on the old version Jax)
• the newer version of Jax was faster for both the compiled and non-compiled example.
I ran 1D logistic regression (like in the readme) but with 2 chains and 3K samples:
Old jax version: Non-compiled: • warmup: 550 samples/sec • sampling: 3800 samples/sec
Compiled: • warmup: 1.9sec • sampling: 1.7 sec
New jax verson (0.2.6): Non-compiled: • warmup: 720 samples/sec • sampling: 5950 samples/sec
Complied: • warmup: 1.5sec • sampling: 1.3 sec
I ran these timings a few times and their variance was small; so the newer jax library definitely sped the sampler up in this example. I don't know if it will make a difference for more complicated examples though!
Nice! Merging.
This fixes the 2 bugs with the older version of Jax (#36 and #58).
I then unpinned the old version of Jax. When I install everything from scratch and run the tests it passes.
Note: I also removed the static_argnums in
integtrators.py
which setsstep_size
(a float) as static. This was also causing a "non-hashable type" error for some reason. In any case, I don't think jitting floats adds a performance penalty (from playing around with Jax generally; I didn't test performance in this particular case).