Add first positive method

UM-ARM-Lab · Jun 7, 2024 · 6312d50 · 6312d50
1 parent f23f15b
commit 6312d50
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/arm_pytorch_utilities/tensor_utils.py b/src/arm_pytorch_utilities/tensor_utils.py
@@ -116,3 +116,7 @@ def ensure_diagonal(Q, dim):
             raise RuntimeError("Expect {} sized diagonal vector but given {}".format(dim, Q.shape[0]))
         Q = torch.diag(Q)
     return Q
+
+def first_positive(x, dim=0):
+    nonz = (x > 0)
+    return ((nonz.cumsum(dim) == 1) & nonz).max(dim)