Rax

What's Changed
* Expose `normalize_probabilities` as a good normalization for `SoftmaxLoss`.
* Remove use of `initial` argument to `jax.nn.softmax` and `jax.nn.log_softmax`
* Drop python 3.8 checks and add python 3.11 checks.
* Changes in lambda weights to reduce boilerplate and add new options.
* Fix pytype and clean up types across codebase.
* Minor typo fixes in documentation.

**Full Changelog**: https://github.com/google/rax/compare/v0.3.0...v0.4.0

What's Changed
* New losses: [`rax.pairwise_soft_zero_one_loss`](https://github.com/google/rax/blob/2d50719937fc4a3caa9967b5a5250f7a980014de/rax/_src/losses.py#L607) and [`rax.pairwise_qr_loss`](https://github.com/google/rax/blob/2d50719937fc4a3caa9967b5a5250f7a980014de/rax/_src/losses.py#L838).
* Added lambdaweight functionality for all pairwise losses (`rax.*_lambdaweight`).
* Added segmented data support for all losses and metrics via the `segments=` keyword argument.
* New example: [Segmentation](https://github.com/google/rax/tree/2d50719937fc4a3caa9967b5a5250f7a980014de/examples/segmentation).

**Full Changelog**: https://github.com/google/rax/compare/v0.2.0...v0.3.0

What's Changed
* New losses: [`rax.poly1_softmax_loss`](https://github.com/google/rax/blob/b5e5fe8702428ba128de84780a0da2e7ade2f401/rax/_src/losses.py#L122) and [`rax.unique_softmax_loss`](https://github.com/google/rax/blob/b5e5fe8702428ba128de84780a0da2e7ade2f401/rax/_src/losses.py#L190).
* New example: [T5X integration](https://github.com/google/rax/tree/b5e5fe8702428ba128de84780a0da2e7ade2f401/examples/t5x).
* Removed dependencies' version upper-bounds.
* Minor documentation fixes.

**Full Changelog**: https://github.com/google/rax/compare/v0.1.0...v0.2.0

This is the initial release of Rax: a Learning-to-Rank (LTR) library built on top of JAX. It includes the following functionality:

* Ranking losses ([`rax.*_loss`](https://rax.readthedocs.io/en/latest/api.html#)).
* Ranking metrics ([`rax.*_metric`](https://rax.readthedocs.io/en/latest/api.html#module-rax._src.metrics)).
* Function transformations ([`rax.*_t12n`](https://rax.readthedocs.io/en/latest/api.html#module-rax._src.t12n)).