From 6312d5094f71a646041fa2994663a19f052ee857 Mon Sep 17 00:00:00 2001 From: Sheng Zhong Date: Fri, 7 Jun 2024 17:53:02 -0400 Subject: [PATCH] Add first positive method --- src/arm_pytorch_utilities/tensor_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/arm_pytorch_utilities/tensor_utils.py b/src/arm_pytorch_utilities/tensor_utils.py index 89f71e8..25d7aa2 100644 --- 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)