From 3991838378ccbd55cb97f031f02fc3eea54cbcac Mon Sep 17 00:00:00 2001
From: Rax Developers <rax-dev@google.com>
Date: Wed, 26 Jun 2024 11:53:17 -0700
Subject: [PATCH] Add safe_reduce to rax.utils for loss/metric implementations
 outside of rax.

PiperOrigin-RevId: 647028192
---
 rax/utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/rax/utils.py b/rax/utils.py
index 4fd1859..ccc2fa2 100644
--- 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",
 ]