Align documentation with up-to-date code.

PiperOrigin-RevId: 649803828
google · Jul 6, 2024 · df63da5 · df63da5
1 parent 3991838
commit df63da5
Showing 1 changed file with 7 additions and 4 deletions.
diff --git 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
 ==========