Fix exp function

Author: RAMitchell

legateboost/metrics.py

@@ -343,8 +343,8 @@ def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         f = cn.log(pred) * (K - self.one)  # undo softmax
         y_k = cn.full((y.size, K), -self.one / (K - self.one))
 
-        set_col_by_idx(y_k, y.astype(cn.int32), self.one)
-        # y_k[cn.arange(y.size), y.astype(cn.int32)] = 1.0
+        # y_k[cn.arange(y.size), y.astype(cn.int32)] = self.one
+        y_k = set_col_by_idx(y_k, y.astype(cn.int32), self.one)
 
         exp = cn.exp(-1 / K * cn.sum(y_k * f, axis=1))
         return (exp * w).sum() / w.sum()

legateboost/objectives.py

@@ -598,7 +598,7 @@ def gradient(self, y: cn.ndarray, pred: cn.ndarray) -> GradPair:
         f = cn.log(pred) * (K - 1)  # undo softmax
         y_k = cn.full((y.size, K), -self.one / (K - self.one))
         labels = y.astype(cn.int32).squeeze()
-        set_col_by_idx(y_k, labels, self.one)
+        y_k = set_col_by_idx(y_k, labels, self.one)
         # y_k[cn.arange(y.size), labels] = 1.0
         exp = cn.exp(-1 / K * cn.sum(y_k * f, axis=1))
 

legateboost/utils.py

@@ -105,7 +105,7 @@ def pick_col_by_idx(a: cn.ndarray, b: cn.ndarray) -> cn.ndarray:
     return result.sum(axis=1)
 
 
-def set_col_by_idx(a: cn.ndarray, b: cn.ndarray, delta: float) -> None:
+def set_col_by_idx(a: cn.ndarray, b: cn.ndarray, delta: float) -> cn.ndarray:
     """Alternative implementation for a[cn.arange(b.size), b] = delta."""
 
     assert a.ndim == 2
@@ -116,7 +116,7 @@ def set_col_by_idx(a: cn.ndarray, b: cn.ndarray, delta: float) -> None:
     bools = b[:, cn.newaxis] == range[cn.newaxis, :]
     a -= a * bools
     a += delta * bools
-    return
+    return a
 
 
 def mod_col_by_idx(a: cn.ndarray, b: cn.ndarray, delta: float) -> None: