Distrax

What's Changed
* Implement TFP's base measure interface for distrax distributions and bijectors. by copybara-service in https://github.com/deepmind/distrax/pull/104
* Add a `matrix` property to affine bijectors that returns a matrix representing their linear part. by copybara-service in https://github.com/deepmind/distrax/pull/114
* Fix categorical KL computation when the second distribution has zero probability on entries where the first distribution also has zero probability. by copybara-service in https://github.com/deepmind/distrax/pull/113
* Implement the `DiagAffine` bijector. by copybara-service in https://github.com/deepmind/distrax/pull/115
* Implement the `MultivariateNormalFromBijector` distribution. by copybara-service in https://github.com/deepmind/distrax/pull/116
* Implement the `MultivariateNormalDiag` distribution as an instance of a `MultivariateNormalFromBijector`. by copybara-service in https://github.com/deepmind/distrax/pull/117
* Implement the `MultivariateNormalTri` distribution. by copybara-service in https://github.com/deepmind/distrax/pull/118
* Implement the `MultivariateNormalDiagPlusLowRank` distribution. by copybara-service in https://github.com/deepmind/distrax/pull/119
* Add tests for KL computations between different Multivariate Gaussian distributions. by copybara-service in https://github.com/deepmind/distrax/pull/123
* Demote `TriangularAffine` to `TriangularLinear`. by copybara-service in https://github.com/deepmind/distrax/pull/125
* Add tests for `same_as` method to increase test coverage. by copybara-service in https://github.com/deepmind/distrax/pull/126
* Demote `DiagAffine` to `DiagLinear`. by copybara-service in https://github.com/deepmind/distrax/pull/127
* Demote `DiagPlusLowRankAffine` to `DiagPlusLowRankLinear`. by copybara-service in https://github.com/deepmind/distrax/pull/128
* Updated 2-distributions tests for the `MultivariateNormalDiagPlusLowRank`. by copybara-service in https://github.com/deepmind/distrax/pull/131
* Adjust test tolerances. by copybara-service in https://github.com/deepmind/distrax/pull/133
* Expose the full-covariance multivariate normal distributions, and upgrade version. by copybara-service in https://github.com/deepmind/distrax/pull/134


**Full Changelog**: https://github.com/deepmind/distrax/compare/v0.1.1...v0.1.2

What's Changed
* Add an absolute tolerance in `multinomial_test` to unblock an XLA optimization. by copybara-service in https://github.com/deepmind/distrax/pull/74
* Change the default `dtype` of discrete distributions to unqualified `int`. by copybara-service in https://github.com/deepmind/distrax/pull/77
* Add a straight-through gradient wrapper method. by copybara-service in https://github.com/deepmind/distrax/pull/70
* Expose `straight_through_wrapper`. by copybara-service in https://github.com/deepmind/distrax/pull/82
* Update Distrax citation. by copybara-service in https://github.com/deepmind/distrax/pull/83
* Ensure forward compatibility with Chex. by copybara-service in https://github.com/deepmind/distrax/pull/84
* Relax test tolerances in `multinomial_test`. by copybara-service in https://github.com/deepmind/distrax/pull/85
* Update requirements and allow new versions of JAX. by copybara-service in https://github.com/deepmind/distrax/pull/87
* Implement a triangular affine bijector, with tests. by copybara-service in https://github.com/deepmind/distrax/pull/88
* Implement `LowerUpperTriangularAffine` as a composition of two `TriangularAffine`s. by copybara-service in https://github.com/deepmind/distrax/pull/89
* Reduce testing time by disabling JAX optimizations. by copybara-service in https://github.com/deepmind/distrax/pull/94
* Remove the old `venv` directory before testing the package. by copybara-service in https://github.com/deepmind/distrax/pull/95
* Initial Gumbel distribution and bijector. by kashif in https://github.com/deepmind/distrax/pull/36
* Add the `Gumbel` distribution. by copybara-service in https://github.com/deepmind/distrax/pull/96
* Delete unnecessary `cdf` definition. by copybara-service in https://github.com/deepmind/distrax/pull/98
* Update the tests for the `Tanh` and `Sigmoid` bijectors. by copybara-service in https://github.com/deepmind/distrax/pull/97
* Delete stale "pylint: disable". by copybara-service in https://github.com/deepmind/distrax/pull/99
* Added the `GumbelCDF` bijector originally developed in https://github.com/deepmind/distrax/pull/36. by copybara-service in https://github.com/deepmind/distrax/pull/100
* Add an affine bijector whose weight matrix is a low-rank perturbation of a diagonal matrix. by copybara-service in https://github.com/deepmind/distrax/pull/90
* Perform validity checking of probability distribution when sampling from `Categorical` / `OneHotCategorical`, and return -1 instead of an invalid sample if normalized probability distribution is invalid. by copybara-service in https://github.com/deepmind/distrax/pull/105
* Use a safer `log_prob` in KL divergence between categoricals. by copybara-service in https://github.com/deepmind/distrax/pull/102
* Clarify behaviour of the Transformed distribution in the docstring. by copybara-service in https://github.com/deepmind/distrax/pull/109


**Full Changelog**: https://github.com/deepmind/distrax/compare/v0.1.0...v0.1.1

Support for Python 3.6 has been dropped as per [JAX deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html). Please upgrade to a supported Python version.

**It is the latest version compatible with Python 3.6. See [deepmind/optax222](https://github.com/deepmind/optax/issues/222) for more details.**

[Full Changelog](https://github.com/deepmind/distrax/compare/v0.0.2...HEAD)

**Closed issues:**

- ImportError: cannot import name 'partial' from 'jax.util' [\60](https://github.com/deepmind/distrax/issues/60)
- Problem with shapes in simple transformed distributions [\42](https://github.com/deepmind/distrax/issues/42)
- Issue with jax version in requirements.txt [\38](https://github.com/deepmind/distrax/issues/38)
- Duplication/Forking vs Collaboration [\35](https://github.com/deepmind/distrax/issues/35)
- Tanh numerical instability [\7](https://github.com/deepmind/distrax/issues/7)

**Merged pull requests:**

- Temporarily disable py3.9 tests. [\69](https://github.com/deepmind/distrax/pull/69) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Minor updates to docstring of bijectors. [\68](https://github.com/deepmind/distrax/pull/68) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Minor updates to slicing tests. [\67](https://github.com/deepmind/distrax/pull/67) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add note for `Tanh` and `Sigmoid`. [\66](https://github.com/deepmind/distrax/pull/66) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Support slicing on more distributions. [\65](https://github.com/deepmind/distrax/pull/65) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- \[JAX\] Update JAX users in preparation for a change that makes iteration over a JAX array return JAX arrays, instead of NumPy arrays. [\63](https://github.com/deepmind/distrax/pull/63) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix typo in `PRNGKey`. [\62](https://github.com/deepmind/distrax/pull/62) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix import errors & freeze the latest compatible JAX version. Fixes \60 [\61](https://github.com/deepmind/distrax/pull/61) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add slice to more distrax distributions. [\59](https://github.com/deepmind/distrax/pull/59) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add slice operation to Distrax Distribution and implement for TFP, MVN and Categorical. [\58](https://github.com/deepmind/distrax/pull/58) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Update tests.yml workflow. [\57](https://github.com/deepmind/distrax/pull/57) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add test.sh for launching CI tests on a local machine. [\56](https://github.com/deepmind/distrax/pull/56) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Use TFP nightly builds in CI tests. [\55](https://github.com/deepmind/distrax/pull/55) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix type of `event_shape_tensor` in Distrax-to-TFP distribution adapter. [\54](https://github.com/deepmind/distrax/pull/54) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Uncomment test cases. [\53](https://github.com/deepmind/distrax/pull/53) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix breakage in the TFP-to-Distrax distribution adapter. [\52](https://github.com/deepmind/distrax/pull/52) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Disable `use-list-literal` refactoring lint tests. [\51](https://github.com/deepmind/distrax/pull/51) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Remove `use-dict-literal` from the set of refactoring lint tests. [\50](https://github.com/deepmind/distrax/pull/50) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Adjust tolerance locally to fix failing test. [\48](https://github.com/deepmind/distrax/pull/48) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Change explicit string formats to f-strings to mitigate github pylint failures. [\47](https://github.com/deepmind/distrax/pull/47) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Replace uses of jax.partial and jax.util.partial with functools.partial, in preparation for removing jax.partial and jax.util.partial. [\46](https://github.com/deepmind/distrax/pull/46) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Remove users of jax.api.\* symbols, in preparation for removing the deprecated jax.api name. [\44](https://github.com/deepmind/distrax/pull/44) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Implement a numerically stable version of the quadratic formula used in the rational quadratic spline bijector. [\43](https://github.com/deepmind/distrax/pull/43) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Distrax: Enable Categorical to be used with a single element. [\41](https://github.com/deepmind/distrax/pull/41) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Replace use of deprecated jax test utility [\40](https://github.com/deepmind/distrax/pull/40) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- initial gamma distribution [\8](https://github.com/deepmind/distrax/pull/8) ([kashif](https://github.com/kashif))



\* *This Changelog was automatically generated by [github_changelog_generator](https://github.com/github-changelog-generator/github-changelog-generator)*

[Full Changelog](https://github.com/deepmind/distrax/compare/v0.0.1...HEAD)

Resolves [the issue with requirements in v0.0.1](https://github.com/deepmind/distrax/issues/38).

**Merged pull requests:**

- Add Hidden Markov Model utility and example. [\37](https://github.com/deepmind/distrax/pull/37) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Broadcast tensors manually when sampling from certain distributions. [\31](https://github.com/deepmind/distrax/pull/31) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix `experimental_shard_axis_names`. [\30](https://github.com/deepmind/distrax/pull/30) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Implement required batch shape methods for TFP compatible bijectors. [\29](https://github.com/deepmind/distrax/pull/29) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- TFP compatibility fix. [\26](https://github.com/deepmind/distrax/pull/26) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Update `requirements.txt`. [\25](https://github.com/deepmind/distrax/pull/25) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Modify the implementation of the forward and inverse methods for the rational quadratic spline to improve speed when using TPUs. [\24](https://github.com/deepmind/distrax/pull/24) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Reduce number of samples in `log_stddev_normal_test.py`. [\22](https://github.com/deepmind/distrax/pull/22) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add more tests in `tfp_compatible_bijector`. [\20](https://github.com/deepmind/distrax/pull/20) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Increase tolerance in rational-quadratic spline test. [\19](https://github.com/deepmind/distrax/pull/19) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Allow the inner bijector to act on arrays rather than just scalars. [\16](https://github.com/deepmind/distrax/pull/16) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix a bug in computing the forward and inverse values at the edges of the range in RationalQuadraticSpline. Numerically x and y values at the edges are the same so it doesn't cause any issues, but it was a conceptual bug. [\14](https://github.com/deepmind/distrax/pull/14) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Expose math.multiply\_no\_nan [\13](https://github.com/deepmind/distrax/pull/13) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Implemented KL divergence for the Transformed distribution. [\12](https://github.com/deepmind/distrax/pull/12) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Updated `flow.py`. [\11](https://github.com/deepmind/distrax/pull/11) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Updated `flow.py`. [\10](https://github.com/deepmind/distrax/pull/10) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add installation instructions. [\2](https://github.com/deepmind/distrax/pull/2) ([copybara-service[bot]](https://github.com/apps/copybara-service))



\* *This Changelog was automatically generated by [github_changelog_generator](https://github.com/github-changelog-generator/github-changelog-generator)*

Changelog

This is a first GitHub release of Distrax!