rax
Learning-to-Rank using JAX.
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires Python 3.9 or later; JAX installation must be completed first per JAX documentation.
- Low friction installation with three runtime dependencies (jax, jaxlib, typing_extensions).
- Repository is actively maintained with recent commits and no known vulnerabilities.
License · maintenance · safety
Apache 2.0 (permissive) — Apache 2.0 permissive license allows commercial and private use with minimal restrictions; suitable for most projects.
last release 2025-01-03 (588 days) · last repo commit 2026-08-06 · 340 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 146,035 downloads/mo, #11,110 on PyPI
Alternatives
Verify before relying
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)- Whether the package is suitable for production ranking systems or primarily research-focused.
- Performance characteristics when scaling to large ranking lists or batch sizes.
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 on this page. Verify before relying on it.
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
rax on PyPI
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.
Requires Python 3.9 or later; JAX installation must be completed first per JAX documentation.
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)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 the current Python release >=3.9 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 3 packagesjaxjaxlibtyping_extensions |
| Maintenance | Actively maintained 588 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 146,035 / month, #11,110 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 4 - BetaIntended Audience :: DevelopersIntended Audience :: Science/ResearchLicense :: OSI Approved :: Apache Software LicenseOperating System :: OS IndependentProgramming Language :: Python :: 3Topic :: Scientific/Engineering :: Artificial IntelligenceTopic :: Software Development :: Libraries :: Python Modules |
Evidence: rax-0.4.0-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “learning to rank jax”
- raxRax provides ranking losses and metrics for learning-to-rank problems…
- chexChex provides utilities for writing reliable JAX code, including…
- cluCLU provides utilities and abstractions for writing machine learning…
Give your agent the search over MCP, or paste the wish link into any chat.
More Python Modules packages
Converts domain names between Unicode and ASCII-compatible encoding (Punycode) according to IDNA 2008 and Unicode Technical Standard 46, with security validation and broader script coverage than the standard library.
Install it if you work with internationalized domain names, need to validate domains, or use HTTP clients that depend on it transitively.
Setuptools is a Python build backend and package management tool that handles building, distributing, and installing Python packages, including support for C/C++ extension modules.
PyYAML parses and emits YAML 1.1 data format, enabling serialization and deserialization of configuration files and Python objects to and from human-readable YAML text.
Pydantic validates Python data structures against type hints, coercing and checking input at runtime to ensure it matches a declared schema.
Provides reusable metadata objects for use with PEP-593 `typing.Annotated` to express common constraints like bounds, collection sizes, and predicates on types.
Install it if you use or build libraries that need to express type constraints in a standardized, inspectable way—or if you want to annotate your own types with…
Provides runtime tools to inspect and introspect Python type annotations, enabling programmatic examination of type hints at execution time.
See also drjax · google-metrax · jraph · pytorch-metric-learning · ranx · clu · dm-haiku · jax · equinox · ir-measures