Add a couple of options to learn-to-rank losses and metrics

PiperOrigin-RevId: 671713845
google · Sep 6, 2024 · 49256cd · 49256cd
1 parent df63da5
commit 49256cd
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/rax/_src/losses.py b/rax/_src/losses.py
@@ -97,7 +97,7 @@ def softmax_loss(
     weights: An optional ``[..., list_size]``-:class:`~jax.Array`, indicating
       the weight for each item.
     label_fn: A label function that maps labels to probabilities. Default keeps
-      labels as-is.
+      labels as-is. See rax.utils.normalize_probabilities for an example.
     reduce_fn: An optional function that reduces the loss values. Can be
       :func:`jax.numpy.sum` or :func:`jax.numpy.mean`. If ``None``, no reduction
       is performed.

diff --git a/rax/utils.py b/rax/utils.py
@@ -19,6 +19,7 @@
 from rax._src.utils import approx_ranks
 from rax._src.utils import compute_pairs
 from rax._src.utils import cutoff
+from rax._src.utils import normalize_probabilities
 from rax._src.utils import ranks
 from rax._src.utils import safe_reduce
 
@@ -27,6 +28,7 @@
     "approx_ranks",
     "compute_pairs",
     "cutoff",
+    "normalize_probabilities",
     "pairwise_loss",
     "ranks",
     "safe_reduce",