---
id: rax
version: "0.4.0"
license: Apache 2.0
license_treatment: permissive
maintenance: active
---

# rax — Learning-to-Rank using JAX.

License: permissive · Maintenance: active · Downloads: 146.0K/mo

## What it is and what it does

Rax is a Learning-to-Rank library built on JAX that solves the problem of optimizing the relative order of items in a list given a context. Unlike classification or regression, ranking problems care about correctness of ordering rather than absolute predictions. Rax provides off-the-shelf ranking losses (softmax_loss, pairwise_logistic_loss), ranking metrics (mrr_metric, ndcg_metric), and transformations (approx_t12n, gumbel_t12n) that compose with standard JAX transformations.

The library is designed for search, recommendation, question-answering, and dialogue systems where you need to score and rank items. All losses and metrics are purely functional and differentiable, allowing you to use them directly in neural network training pipelines. You can also transform non-differentiable ranking metrics into differentiable losses, enabling gradient-based optimization of metrics like NDCG.

Use it for:

- Ranking documents in search engines by relevance to a query using differentiable ranking losses.
- Building recommendation systems that rank items for a user and optimize for ranking metrics like MRR.
- Training question-answering systems to rank candidate answers by relevance.
- Optimizing dialogue systems to select the best response from a list of candidates.
- Implementing pairwise or listwise ranking losses in JAX-based machine learning pipelines.

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

Rax provides ranking losses and metrics for learning-to-rank problems in JAX, enabling optimization of relative ordering for lists of items given a context.

Yes. Rax fills a clear gap for learning-to-rank in the JAX ecosystem with low install friction, active maintenance, no known vulnerabilities, and permissive licensing. Install it if you're building ranking systems in JAX or need to optimize ranking metrics directly.

## Install

pip install rax
uv add rax
poetry add rax

## Installing rax

Before you install: Low friction installation with three runtime dependencies (jax, jaxlib, typing_extensions). Repository is actively maintained with recent commits and no known vulnerabilities.

License in practice: Apache 2.0 permissive license allows commercial and private use with minimal restrictions; suitable for most projects.

Quickstart:

pip install rax

import jax.numpy as jnp
import rax

scores = jnp.array([2.2, -1.3, 5.4])
labels = jnp.array([1.0, 0.0, 0.0])
rax.ndcg_metric(scores, labels)

Requires Python 3.9 or later; JAX installation must be completed first per JAX documentation.

Verify before relying:

- Whether the package is suitable for production ranking systems or primarily research-focused.
- Performance characteristics when scaling to large ranking lists or batch sizes.

## Package facts

- License: Apache 2.0 (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 146.0K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

learning to rank jax, ranking losses metrics, learning-to-rank library, jax ranking optimization, pairwise ranking loss, ndcg metric computation, search ranking neural networks, learning-to-rank, jax-ecosystem, ranking-metrics

[View on SkillFed](https://skillfed.io/packages/rax) · [View on PyPI](https://pypi.org/project/rax/)