Numpyro: standard_gamma not caching compiled function

[neerajprad]

I would have expected the second call to Beta.sample() to be much faster due to 2 reasons - no compilation cost for standard_gamma, and faster samples from the compiled kernel.

This seems to not be the case.

In [4]: %timeit Beta(1.1, 1.1).sample(PRNGKey(1), (1000,))
2.24 s ± 15.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [5]: %timeit Beta(1.1, 1.1).sample(PRNGKey(1), (1000,))
2.3 s ± 46.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[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]

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]

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.

[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: https://github.com/google/jax/issues/695. :)

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