standard_gamma not caching compiled function about numpyro

[fehiepsi]

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

[fehiepsi]

@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. :)

[fehiepsi]

This is bad. I believe this is a regression because it happens for standard_gamma too.

[neerajprad]

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.

[fehiepsi]

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. >"<

[neerajprad]

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.

[fehiepsi]

The fix is landed in the new version of JAX. Hopefully our tests will be faster now.

[neerajprad]

@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!