Closed sidravi1 closed 3 years ago
I think I'm mostly there with shape. Will run a bunch of tests tomorrow and if it's all kosher, I'll put it up for another review.
Added some test cases. I think shape stuff is looking ok now.
Batching when using jax.scipy's multivariate_normal's logp gives inf
s. Not sure if I'm doing something stupid. I raised this.
Patched, now the tests pass!
Made the fixes!
Re: the jax bug. The tests are not affected since they are only checking shape. The jax bug leads to inf
being returned but the shape is correct.
Btw... Do you want me to squash all my merges so the commit history doesn't get messy?
Re: the jax bug. The tests are not affected since they are only checking shape. The jax bug leads to
inf
being returned but the shape is correct.
Ok, hopefully it will be fixed before the first release.
Btw... Do you want me to squash all my merges so the commit history doesn't get messy?
Yes please!
Looks good to me, merging. Great job!
Addressed one item in #65
Overview
Add MvNormal distribution. Uses
jax.scipy.multivariate_normal
forlogpdf
andjax.random.multivariate_normal
for sampling.Quick tests
Not sure the batching is broadcasting correctly though:
Pending
Proper test cases Fixing shape issues