Comments (8)
Yup, I can confirm this is a regression of JAX. This happens for other samplers too: e.g. random.normal
.
random.normal(random.PRNGKey(0), (1000,))
Jax 0.1.24 took 1ms
while Jax 0.1.25 took more than 100ms
cc @mattjj
from numpyro.
@neerajprad It seems that LKJ does not perform well in CPU w.r.t. GPU is because of vector_to_tril_matrix
function, not because of gamma sampler. Most of 5s in the benchmark is spent for this function. I reported the issue here: google/jax#695. :)
from numpyro.
This is bad. I believe this is a regression because it happens for standard_gamma
too.
from numpyro.
I believe this is a regression because it happens for standard_gamma too.
I noticed this because test_distributions.py
seems to be unusually slow now (earlier it was only initial calls to some tests during which standard_gamma
would be compiled for a few shapes), so I think this is a regression too.
from numpyro.
It seems a bit weird that if I wrap these samplers in a dummy function
def f():
return random.normal(random.PRNGKey(0), (1000,))
then it is fast again.
This applies for standard_gamma
too but does not apply for our distribution class. >"<
from numpyro.
This applies for standard_gamma too but does not apply for our distribution class. >"<
Hmm..maybe the cache key is different now and for some reason we are generating a new key with our distribution classes and missing the cache. I'll take a look at it too.
from numpyro.
The fix is landed in the new version of JAX. Hopefully our tests will be faster now.
from numpyro.
@neerajprad It seems that LKJ does not perform well in CPU w.r.t. GPU is because of vector_to_tril_matrix function, not because of gamma sampler.
Whoa..that's great to know. Thanks for digging into this, and pinning it down!
from numpyro.
Related Issues (20)
- [FR] Truncated Power Law distribution
- An auto guide's `_unpack_latent` and `_unpack_latent._inverse` don't use produce the same order HOT 1
- [FR] Utility for joint distributions HOT 14
- How can I gibbs before HMC/NUTS? HOT 8
- Large potential energy while using `HMCGibbs` at the initial stage HOT 3
- Inference Test Failing HOT 2
- Figure in AR2 example is not reproducible
- Got Problems When Computing Log Likelihoods in a Scan-Based VAR Model HOT 2
- Inference for `gaussian_hmm` is broken on latest jax version (0.4.30) HOT 1
- autocorrelation function HOT 5
- saving render_model() output to the desired file path HOT 5
- Object oriented wrapper API HOT 5
- Stress test utility for numpyro? HOT 2
- Samples are outside the support for DiscreteUniform distribution HOT 3
- Crash when using `TruncatedNormal` in `parallel` MCMC, but not in `sequential` MCMC HOT 5
- `nuts.get_extra_fields()["num_steps"]=0` after warmup HOT 4
- Trouble Initializing Parameters with Highly Degenerate Data HOT 1
- Predictive: better docs for setting obs to `None` HOT 4
- Error in model using CircularReparam when trying to use Predictive HOT 1
- Error when using obs_mask and predictive with different input shape HOT 2
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from numpyro.