# 🦖 **Rax**: Learning-to-Rank using JAX
[](https://rax.readthedocs.io/en/latest/?badge=latest)
[](https://pypi.org/project/rax/)
[](https://github.com/google/rax/blob/main/LICENSE)
**Rax** is a Learning-to-Rank library written in JAX. Rax provides off-the-shelf
implementations of ranking losses and metrics to be used with JAX. It provides
the following functionality:
- Ranking losses (`rax.*_loss`): `rax.softmax_loss`,
`rax.pairwise_logistic_loss`, ...
- Ranking metrics (`rax.*_metric`): `rax.mrr_metric`, `rax.ndcg_metric`, ...
- Transformations (`rax.*_t12n`): `rax.approx_t12n`, `rax.gumbel_t12n`, ...
## Ranking
A ranking problem is different from traditional classification/regression
problems in that its objective is to optimize for the correctness of the
**relative order** of a **list of examples** (e.g., documents) for a given
context (e.g., a query). **Rax** provides support for ranking problems within
the JAX ecosystem. It can be used in, but is not limited to, the following
applications:
- **Search**: ranking a list of documents with respect to a query.
- **Recommendation**: ranking a list of items given a user as context.
- **Question Answering**: finding the best answer from a list of candidates.
- **Dialogue System**: finding the best response from a list of responses.
## Synopsis
In a nutshell, given the scores and labels for a list of items, Rax can compute
various ranking losses and metrics:
```python
import jax.numpy as jnp
import rax
scores = jnp.array([2.2, -1.3, 5.4]) # output of a model.
labels = jnp.array([1.0, 0.0, 0.0]) # indicates doc 1 is relevant.
rax.ndcg_metric(scores, labels) # computes a ranking metric.
# 0.63092977
rax.pairwise_hinge_loss(scores, labels) # computes a ranking loss.
# 2.1
```
All of the Rax losses and metrics are purely functional and compose well with
standard JAX transformations. Additionally, Rax provides ranking-specific
transformations so you can build new ranking losses. An example is
`rax.approx_t12n`, which can be used to transform any (non-differentiable)
ranking metric into a differentiable loss. For example:
```python
loss_fn = rax.approx_t12n(rax.ndcg_metric)
loss_fn(scores, labels) # differentiable approx ndcg loss.
# -0.63282484
jax.grad(loss_fn)(scores, labels) # computes gradients w.r.t. scores.
# [-0.01276882 0.00549765 0.00727116]
```
## Installation
See https://github.com/google/jax#installation for instructions on installing JAX.
We suggest installing the latest stable version of Rax by running:
```
$ pip install rax
```
## Examples
See the `examples/` directory for complete examples on how to use Rax.
## Citing Rax
If you use Rax, please consider citing our
[paper](https://research.google/pubs/pub51453/):
```
@inproceedings{jagerman2022rax,
title = {Rax: Composable Learning-to-Rank using JAX},
author = {Rolf Jagerman and Xuanhui Wang and Honglei Zhuang and Zhen Qin and
Michael Bendersky and Marc Najork},
year = {2022},
booktitle = {Proceedings of the 28th ACM SIGKDD International Conference on Knowledge Discovery \& Data Mining}
}
```
Raw data
{
"_id": null,
"home_page": "https://github.com/google/rax",
"name": "rax",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": null,
"keywords": "learning-to-rank jax ranking",
"author": "Google",
"author_email": "rax-dev@google.com",
"download_url": "https://files.pythonhosted.org/packages/64/4d/38472764ff89fecad4f4eb81cabf1c29c2cab02a843baa979c7797a0db22/rax-0.4.0.tar.gz",
"platform": null,
"description": "# \ud83e\udd96 **Rax**: Learning-to-Rank using JAX\n\n[](https://rax.readthedocs.io/en/latest/?badge=latest)\n[](https://pypi.org/project/rax/)\n[](https://github.com/google/rax/blob/main/LICENSE)\n\n**Rax** is a Learning-to-Rank library written in JAX. Rax provides off-the-shelf\nimplementations of ranking losses and metrics to be used with JAX. It provides\nthe following functionality:\n\n- Ranking losses (`rax.*_loss`): `rax.softmax_loss`,\n `rax.pairwise_logistic_loss`, ...\n- Ranking metrics (`rax.*_metric`): `rax.mrr_metric`, `rax.ndcg_metric`, ...\n- Transformations (`rax.*_t12n`): `rax.approx_t12n`, `rax.gumbel_t12n`, ...\n\n## Ranking\n\nA ranking problem is different from traditional classification/regression\nproblems in that its objective is to optimize for the correctness of the\n**relative order** of a **list of examples** (e.g., documents) for a given\ncontext (e.g., a query). **Rax** provides support for ranking problems within\nthe JAX ecosystem. It can be used in, but is not limited to, the following\napplications:\n\n- **Search**: ranking a list of documents with respect to a query.\n- **Recommendation**: ranking a list of items given a user as context.\n- **Question Answering**: finding the best answer from a list of candidates.\n- **Dialogue System**: finding the best response from a list of responses.\n\n## Synopsis\n\nIn a nutshell, given the scores and labels for a list of items, Rax can compute\nvarious ranking losses and metrics:\n\n```python\nimport jax.numpy as jnp\nimport rax\n\nscores = jnp.array([2.2, -1.3, 5.4]) # output of a model.\nlabels = jnp.array([1.0, 0.0, 0.0]) # indicates doc 1 is relevant.\n\nrax.ndcg_metric(scores, labels) # computes a ranking metric.\n# 0.63092977\n\nrax.pairwise_hinge_loss(scores, labels) # computes a ranking loss.\n# 2.1\n```\n\nAll of the Rax losses and metrics are purely functional and compose well with\nstandard JAX transformations. Additionally, Rax provides ranking-specific\ntransformations so you can build new ranking losses. An example is\n`rax.approx_t12n`, which can be used to transform any (non-differentiable)\nranking metric into a differentiable loss. For example:\n\n```python\nloss_fn = rax.approx_t12n(rax.ndcg_metric)\nloss_fn(scores, labels) # differentiable approx ndcg loss.\n# -0.63282484\n\njax.grad(loss_fn)(scores, labels) # computes gradients w.r.t. scores.\n# [-0.01276882 0.00549765 0.00727116]\n```\n\n## Installation\n\nSee https://github.com/google/jax#installation for instructions on installing JAX.\n\nWe suggest installing the latest stable version of Rax by running:\n\n```\n$ pip install rax\n```\n\n## Examples\n\nSee the `examples/` directory for complete examples on how to use Rax.\n\n## Citing Rax\n\nIf you use Rax, please consider citing our\n[paper](https://research.google/pubs/pub51453/):\n\n```\n@inproceedings{jagerman2022rax,\n title = {Rax: Composable Learning-to-Rank using JAX},\n author = {Rolf Jagerman and Xuanhui Wang and Honglei Zhuang and Zhen Qin and\n Michael Bendersky and Marc Najork},\n year = {2022},\n booktitle = {Proceedings of the 28th ACM SIGKDD International Conference on Knowledge Discovery \\& Data Mining}\n}\n```\n",
"bugtrack_url": null,
"license": "Apache 2.0",
"summary": "Learning-to-Rank using JAX.",
"version": "0.4.0",
"project_urls": {
"Homepage": "https://github.com/google/rax"
},
"split_keywords": [
"learning-to-rank",
"jax",
"ranking"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "b7d6b274635c0cfb1925fc828dfb6fe7f3acdfda04d41e0bda04389f44836dd4",
"md5": "675a980de923c2929cfb2b020ed65a1c",
"sha256": "9cab65339c1191a3f200d406f9d6e141fa1d6410f696e4cf6a8869d1487a565c"
},
"downloads": -1,
"filename": "rax-0.4.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "675a980de923c2929cfb2b020ed65a1c",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 79883,
"upload_time": "2025-01-03T19:06:40",
"upload_time_iso_8601": "2025-01-03T19:06:40.575871Z",
"url": "https://files.pythonhosted.org/packages/b7/d6/b274635c0cfb1925fc828dfb6fe7f3acdfda04d41e0bda04389f44836dd4/rax-0.4.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "644d38472764ff89fecad4f4eb81cabf1c29c2cab02a843baa979c7797a0db22",
"md5": "f9a7ee3aa85df1c810e62284d716332c",
"sha256": "1d85164c8bbf712566798b6ac04a40a49a471a6dce235692c051ffa70a2589da"
},
"downloads": -1,
"filename": "rax-0.4.0.tar.gz",
"has_sig": false,
"md5_digest": "f9a7ee3aa85df1c810e62284d716332c",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 60445,
"upload_time": "2025-01-03T19:06:42",
"upload_time_iso_8601": "2025-01-03T19:06:42.689912Z",
"url": "https://files.pythonhosted.org/packages/64/4d/38472764ff89fecad4f4eb81cabf1c29c2cab02a843baa979c7797a0db22/rax-0.4.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-01-03 19:06:42",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "google",
"github_project": "rax",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "rax"
}
JSON
