From df63da50f4c1f39a5a55fe295a468a78bb5aa2a3 Mon Sep 17 00:00:00 2001
From: Rolf Jagerman <jagerman@google.com>
Date: Sat, 6 Jul 2024 02:35:35 -0700
Subject: [PATCH] Align documentation with up-to-date code.

PiperOrigin-RevId: 649803828
---
 docs/api.rst | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/docs/api.rst b/docs/api.rst
index d3bc7b4..261ae63 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -94,6 +94,7 @@ Lambdaweights (``rax.*_lambdaweight``)
 Utilities
 =========
 
+.. automodule:: rax._src.utils
 .. currentmodule:: rax.utils
 
 .. autosummary::
@@ -102,11 +103,17 @@ Utilities
     cutoff
     approx_ranks
     approx_cutoff
+    pairwise_loss
+    compute_pairs
+    safe_reduce
 
 .. autofunction:: ranks
 .. autofunction:: cutoff
 .. autofunction:: approx_ranks
 .. autofunction:: approx_cutoff
+.. autofunction:: pairwise_loss
+.. autofunction:: compute_pairs
+.. autofunction:: safe_reduce
 
 Types
 =====
@@ -122,7 +129,6 @@ Types
     MetricFn
     RankFn
     ReduceFn
-    WeightFn
 
 .. autoclass:: CutoffFn
    :members: __call__
@@ -142,9 +148,6 @@ Types
 .. autoclass:: ReduceFn
    :members: __call__
 
-.. autoclass:: WeightFn
-   :members: __call__
-
 References
 ==========