Add safe_reduce to rax.utils for loss/metric implementations outside …

…of rax. PiperOrigin-RevId: 647028192
google · Jun 26, 2024 · 3991838 · 3991838
1 parent 1638551
commit 3991838
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/rax/utils.py b/rax/utils.py
@@ -20,6 +20,7 @@
 from rax._src.utils import compute_pairs
 from rax._src.utils import cutoff
 from rax._src.utils import ranks
+from rax._src.utils import safe_reduce
 
 __all__ = [
     "approx_cutoff",
@@ -28,4 +29,5 @@
     "cutoff",
     "pairwise_loss",
     "ranks",
+    "safe_reduce",
 ]