initial gumbel distribution and bijection

[kashif]

for issue #21

[franrruiz]

Very nice, thank you! I left some (mostly minor) comments. I would also suggest to add a gumbel_cdf_test.py file to test that the bijector behaves as intended.

[kashif]

@franrruiz thanks for the review! I appreciate the time you spend on it and I'll get it fixed.

[franrruiz]

Hi Kashif - Just following up on the status of this CL. It's really nice and I'd like to have it imported to the library. Are there are any outstanding comments for which you'd like to get some input? I can see that there seem to be 2 tests that aren't passing which seem to refer to the number of arguments of the entropy method - that should be easily fixable but let me know if it's actually harder than it looks like.

[kashif]

Ok cool I’ll have a look and fix it up.

[kashif]

@franrruiz I'm really stuck trying to figure out why the tests are failing... would you have a few mins to have a look and help me out? thanks!

[kashif]

thanks! I appreciate it!

[franrruiz]

My pleasure! Just a quick update on why the tests were failing:

1. Most tests failed because of TFP Gumbel. In particular, the tests use both Distrax and TFP distributions to check that they agree. But when creating the TFP distribution, the loc and scale parameters needed to be numpy arrays of type np.float32; otherwise an error was raised.

2. Some other tests failed because the self._batch_size property of the Gumbel distribution was set to None when calling to the super.__init__ method. To fix this, I just moved the line where we set self._batch_size to appear after the call to the __init__ method of the superclass.

3. As a minor point, there was a jnp.lgamma that should have been jax.lax.lgamma.

Besides this, we are planning to make a few more updates to the distribution and bijector before exposing it in the library. It'll be ready soon :)

[franrruiz]

Thank you @kashif for taking the initiative to write this addition to the library!

[kashif]

awesome! I'll take these lessons for the next distribution/transformation I implement! Very useful!