What's Changed
* Fix failing copybara lint errors. by copybara-service in https://github.com/deepmind/rlax/pull/24
* Add tests for clipped_entropy_softmax distribution and fix improperly negated clipped entropy. by copybara-service in https://github.com/deepmind/rlax/pull/21
* Add tests for multivariate_normal_kl_divergence & kl functions in gaussian_diagonal. by copybara-service in https://github.com/deepmind/rlax/pull/22
* Migrate RLax distributions to use distrax. by copybara-service in https://github.com/deepmind/rlax/pull/20
* Re-allow rlax gaussian diagonal to work with scalar sigma. by copybara-service in https://github.com/deepmind/rlax/pull/25
* Fixes bug in kl calculation of gaussian_diagonal by copybara-service in https://github.com/deepmind/rlax/pull/26
* [JAX] Replace uses of deprecated `jax.ops.index_update(x, idx, y)` APIs with their up-to-date, more succinct equivalent `x.at[idx].set(y)`. by copybara-service in https://github.com/deepmind/rlax/pull/33
* [JAX] Increase numerical tolerances of tests in preparation for an XLA:CPU vectorization change. by copybara-service in https://github.com/deepmind/rlax/pull/34
* Fix performance issue in simple DQN example. by copybara-service in https://github.com/deepmind/rlax/pull/32
* Add test.sh for launching CI tests on a local machine. by copybara-service in https://github.com/deepmind/rlax/pull/35
* Iterate over Python range instead jnp.arange. by copybara-service in https://github.com/deepmind/rlax/pull/29
* fix kl argument order for gaussians by akssri-sai in https://github.com/deepmind/rlax/pull/15
* Freeze the latest compatible JAX version. by copybara-service in https://github.com/deepmind/rlax/pull/36
* Internal change. by copybara-service in https://github.com/deepmind/rlax/pull/38
**Full Changelog**: https://github.com/deepmind/rlax/compare/v0.0.4...v0.0.5