Complete tests for potrs

vishwakftw · vishwakftw · commit 551360a3c6d7 · 2018-11-03T19:03:35.000+05:30
diff --git a/aten/src/ATen/Declarations.cwrap b/aten/src/ATen/Declarations.cwrap
@@ -2738,7 +2738,7 @@
       default: U
 ]]
 [[
-  name: _potrs_single
+  name: _th_potrs_single
   cname: potrs
   types:
     - Float
diff --git a/aten/src/ATen/core/aten_interned_strings.h b/aten/src/ATen/core/aten_interned_strings.h
@@ -81,8 +81,6 @@ _(aten, _floor) \
 _(aten, _fused_dropout) \
 _(aten, _ger) \
 _(aten, _gesv_helper) \
-_(aten, _gesv_single) \
-_(aten, _getri_single) \
 _(aten, _indexCopy) \
 _(aten, _indices) \
 _(aten, _inverse_helper) \
@@ -104,7 +102,6 @@ _(aten, _pad_packed_sequence) \
 _(aten, _pdist_backward) \
 _(aten, _pdist_forward) \
 _(aten, _potrs_helper) \
-_(aten, _potrs_single) \
 _(aten, _prod) \
 _(aten, _prodall) \
 _(aten, _range) \
diff --git a/aten/src/ATen/native/BatchLinearAlgebra.cpp b/aten/src/ATen/native/BatchLinearAlgebra.cpp
@@ -269,7 +269,7 @@ Tensor _potrs_helper_cpu(const Tensor& self, const Tensor& A, bool upper) {
 // Supports arbitrary batch dimensions for self and A
 Tensor potrs(const Tensor& self, const Tensor& A, bool upper) {
   if (self.dim() <= 2 && A.dim() <= 2) {
-    return at::_potrs_single(self, A, upper);
+    return at::_th_potrs_single(self, A, upper);
   }
 
   Tensor self_broadcasted, A_broadcasted;
@@ -281,7 +281,7 @@ Tensor& potrs_out(Tensor& result, const Tensor& self, const Tensor& A, bool uppe
   AT_CHECK(self.dim() == 2 && A.dim() == 2,
            "torch.potrs() with the `out` keyword does not support batching. "
            "b.dim() (", self.dim(), ") and A.dim() (", A.dim(), ") must both be 2.");
-  return at::_potrs_single_out(result, self, A, upper);
+  return at::_th_potrs_single_out(result, self, A, upper);
 }
 
 }}  // namespace at::native
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -1557,6 +1557,10 @@ def test_potrs(self):
     def test_potrs_batched(self):
         _TestTorchMixin._test_potrs_batched(self, lambda t: t.cuda())
 
+    @unittest.skipIf(not TEST_MAGMA, "no MAGMA library detected")
+    def test_potrs_batched_dims(self):
+        _TestTorchMixin._test_potrs_batched_dims(self, lambda t: t.cuda())
+
     def test_view(self):
         _TestTorchMixin._test_view(self, lambda t: t.cuda())
 
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -5404,6 +5404,57 @@ def get_cholesky(bmat, upper):
     def test_potrs_batched(self):
         self._test_potrs_batched(self, lambda t: t)
 
+    @staticmethod
+    def _test_potrs_batched_dims(self, cast):
+        if not TEST_NUMPY:
+            return
+
+        from numpy.linalg import solve
+        from common_utils import random_symmetric_pd_matrix
+
+        for upper in [True, False]:
+            # TODO: This function should be replaced after batch potrf is ready
+            def get_cholesky(bmat, upper):
+                n = bmat.size(-1)
+                cholesky = torch.stack([m.cholesky(upper) for m in bmat.reshape(-1, n, n)])
+                return cholesky.reshape_as(bmat)
+
+            # test against numpy.linalg.solve
+            A = cast(random_symmetric_pd_matrix(4, 2, 1, 3))
+            b = cast(torch.randn(2, 1, 3, 4, 6))
+            L = get_cholesky(A, upper)
+            x = torch.potrs(b, L, upper=upper)
+            x_exp = torch.Tensor(solve(A.cpu().numpy(), b.cpu().numpy()))
+            self.assertEqual(x.data, cast(x_exp))
+
+            # broadcasting b
+            A = cast(random_symmetric_pd_matrix(4, 2, 1, 3))
+            b = cast(torch.randn(4, 6))
+            L = get_cholesky(A, upper)
+            x = torch.potrs(b, L, upper=upper)
+            x_exp = torch.Tensor(solve(A.cpu().numpy(), b.cpu().numpy()))
+            self.assertEqual(x.data, cast(x_exp))
+
+            # broadcasting A
+            A = cast(random_symmetric_pd_matrix(4))
+            b = cast(torch.randn(2, 1, 3, 4, 2))
+            L = get_cholesky(A, upper)
+            x = torch.potrs(b, L, upper=upper)
+            x_exp = torch.Tensor(solve(A.cpu().numpy(), b.cpu().numpy()))
+            self.assertEqual(x.data, cast(x_exp))
+
+            # broadcasting both A & b
+            A = cast(random_symmetric_pd_matrix(4, 1, 3, 1))
+            b = cast(torch.randn(2, 1, 3, 4, 5))
+            L = get_cholesky(A, upper)
+            x = torch.potrs(b, L, upper=upper)
+            x_exp = torch.Tensor(solve(A.cpu().numpy(), b.cpu().numpy()))
+            self.assertEqual(x.data, cast(x_exp))
+
+    @skipIfNoLapack
+    def test_potrs_batched_dims(self):
+        self._test_potrs_batched_dims(self, lambda t: t)
+
     @skipIfNoLapack
     def test_potri(self):
         a = torch.Tensor(((6.80, -2.11, 5.66, 5.97, 8.23),
