Seperate v into v+ and v-
@@ -30,7 +30,7 @@ class MulLazyVariable(LazyVariable):
added_diag = None
if len(self.lazy_vars) == 1:
- def closure(rhs_mat):
+ def _closure(rhs_mat):
if_vector = False
if rhs_mat.ndimension() == 1:
rhs_mat = rhs_mat.unsqueeze(1)
@@ -43,9 +43,8 @@ class MulLazyVariable(LazyVariable):
res = res_mul
res = res.squeeze(1) if if_vector else res
return res
- return closure
-
- def closure(rhs_mat):
+ else:
+ def _closure(rhs_mat):
if_vector = False
if rhs_mat.ndimension() == 1:
rhs_mat = rhs_mat.unsqueeze(1)
@@ -80,6 +79,12 @@ class MulLazyVariable(LazyVariable):
res = res.squeeze(1) if if_vector else res
return res
+ def closure(rhs_mat):
+ zeros = torch.zeros(rhs_mat.size())
+ rhs_mat_pos = torch.max(rhs_mat, zeros)
+ rhs_mat_neg = torch.min(rhs_mat, zeros)
+ return _closure(rhs_mat_pos) + _closure(rhs_mat_neg)
+
return closure
def _derivative_quadratic_form_factory(self, *args):
