Registry / ai-ml / rax

rax

JSON →

library0.4.0pypypi✓ verified 82d ago

Learning-to-Rank using JAX. Current version 0.4.0. Development is active but slow; no recent releases in 2024.

ai-ml

pip install rax

github ↗homepage ↗

INSTALL

IMPORT

SIG · RAX

rax

ai-mlpythonv0.4.0

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

PointwiseLoss

✓ from rax import pointwise_mse_loss

✗ from rax import PointwiseLoss

Compute pointwise MSE loss and NDCG metric

import jax.numpy as jnp from rax import PointwiseLoss, NDCG scores = jnp.array([2.0, 1.0, 0.0]) labels = jnp.array([1.0, 0.0, 1.0]) loss = PointwiseLoss('mse')(scores=scores, labels=labels) print(loss) ndcg = NDCG()(scores=scores, labels=labels) print(ndcg)

Debug

Known issues

gotcharax uses JAX's functional transformations (e.g., jax.grad) but loss functions return a scalar. For gradients, you must call jax.grad on the loss function wrapper, not the loss instance directly.

fix

Use `grad_loss = jax.grad(lambda s: loss(scores=s, labels=labels))`

affects: all

breakingIn rax 0.3.0, the import paths changed. All losses and metrics are now importable directly from the `rax` package instead of submodules like `rax.losses`.

fix

Change `from rax.losses import PointwiseLoss` to `from rax import PointwiseLoss`

affects: 0.3.0+

deprecatedThe `one_hot_labels` parameter in some loss functions is deprecated in favor of using `labels` directly with integer or float arrays.

fix

Pass labels as integers (e.g., 0,1,2) instead of one-hot vectors.

affects: 0.4.0

gotchaMetrics like NDCG expect scores and labels to be 2D (batch x list_size). If you pass 1D arrays, you may get incorrect results without error.

fix

Ensure arrays are of shape (batch_size, list_size) using `jnp.expand_dims` or batching.

affects: all

Upgrade

Version history

0.4.0latest on PyPI · released Jan 3, 2025

Audit

Dependencies

jaxrequiredrax is built on JAX

jaxlibrequiredJAX runtime (CPU/GPU)

Agent activity

10 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/google/rax ↗

homepagehttps://github.com/google/rax ↗