Add Gesv to autograd (pytorch#1733)

* Add Gesv to autograd

* Add TODO for backprop through LU

Author: fmassa

test/test_autograd.py

@@ -1419,6 +1419,7 @@ class dont_convert(tuple):
     (Cross, (), ((S, 3), (S, 3))),
     (Cross, (), ((S, 3, S), (S, 3, S), 1), 'dim'),
     (Inverse, (), ((S, S),), '', (), [skipIfNoLapack]),
+    (Gesv, (), ((S, S), (S, S)), '', (), [skipIfNoLapack]),
     (Clone, (), ((S, M, S),)),
     (Squeeze, (), ((S, 1, M, 1),)),
     # TODO: enable neg dim checks
@@ -1552,6 +1553,7 @@ class dont_convert(tuple):
     ('cross', (S, 3), ((S, 3),)),
     ('cross', (S, 3, S), ((S, 3, S), 1), 'dim'),
     ('inverse', (S, S), (), '', (), [skipIfNoLapack]),
+    ('gesv', (S, S), ((S, S),), '', (), [skipIfNoLapack]),
     ('clone', (S, M, S), ()),
     ('eq', (S, S, S), ((S, S, S),)),
     ('ne', (S, S, S), ((S, S, S),)),

torch/autograd/_functions/linalg.py

@@ -84,3 +84,21 @@ def forward(ctx, input):
     def backward(ctx, grad_output):
         inverse, = ctx.saved_variables
         return -torch.mm(inverse.t(), torch.mm(grad_output, inverse.t()))
+
+
+class Gesv(Function):
+
+    @staticmethod
+    def forward(ctx, b, a):
+        # TODO see if one can backprop through LU
+        X, LU = torch.gesv(b, a)
+        ctx.save_for_backward(X, a)
+        ctx.mark_non_differentiable(LU)
+        return X, LU
+
+    @staticmethod
+    def backward(ctx, grad_output, grad_LU=None):
+        X, a = ctx.saved_variables
+        grad_b, _ = torch.gesv(grad_output, a.t())
+        grad_a = -torch.mm(grad_b, X.t())
+        return grad_b, grad_a

torch/autograd/variable.py

@@ -719,6 +719,9 @@ def cross(self, other, dim=-1):
     def inverse(self):
         return Inverse.apply(self)
 
+    def gesv(self, a):
+        return Gesv.apply(self, a)
+
     def multinomial(self, num_samples=1, with_replacement=False):
         return Multinomial(num_samples, with_replacement)(self)