Fix make_node for cpu svd and add infer_shape for svd

theano/gpuarray/tests/test_linalg.py

@@ -292,7 +292,7 @@ class TestMagma(unittest.TestCase):
             mode=mode_with_gpu.including('magma'))
 
         A_val = rand(50, 100)
-        utt.assert_allclose(f_cpu(A_val)[1], f_gpu(A_val)[1])
+        utt.assert_allclose(f_cpu(A_val), f_gpu(A_val)[1])
 
         A_val = rand(100, 50)
-        utt.assert_allclose(f_cpu(A_val)[1], f_gpu(A_val)[1])
+        utt.assert_allclose(f_cpu(A_val), f_gpu(A_val)[1])

theano/tensor/nlinalg.py

@@ -616,21 +616,37 @@ class SVD(Op):
     def make_node(self, x):
         x = as_tensor_variable(x)
         assert x.ndim == 2, "The input of svd function should be a matrix."
-        w = theano.tensor.matrix(dtype=x.dtype)
-        u = theano.tensor.vector(dtype=x.dtype)
-        v = theano.tensor.matrix(dtype=x.dtype)
-        return Apply(self, [x], [w, u, v])
+        s = theano.tensor.vector(dtype=x.dtype)
+        if self.compute_uv:
+            u = theano.tensor.matrix(dtype=x.dtype)
+            vt = theano.tensor.matrix(dtype=x.dtype)
+            return Apply(self, [x], [u, s, vt])
+        else:
+            return Apply(self, [x], [s])
 
     def perform(self, node, inputs, outputs):
         (x,) = inputs
-        (w, u, v) = outputs
         assert x.ndim == 2, "The input of svd function should be a matrix."
         if self.compute_uv:
-            w[0], u[0], v[0] = self._numop(x,
-                                           self.full_matrices,
-                                           self.compute_uv)
+            u, s, vt = outputs
+            u[0], s[0], vt[0] = self._numop(x,
+                                            self.full_matrices,
+                                            self.compute_uv)
+        else:
+            s, = outputs
+            s[0] = self._numop(x, self.full_matrices, self.compute_uv)
+
+    def infer_shape(self, node, shapes):
+        x_shape, = shapes
+        M, N = x_shape
+        K = tensor.minimum(M, N)
+        s_shape = (K, )
+        if self.compute_uv:
+            u_shape = (M, M) if self.full_matrices else (M, K)
+            vt_shape = (N, N) if self.full_matrices else (K, N)
+            return [u_shape, s_shape, vt_shape]
         else:
-            u[0] = self._numop(x, self.full_matrices, self.compute_uv)
+            return [s_shape]
 
 
 def svd(a, full_matrices=1, compute_uv=1):

theano/tensor/tests/test_nlinalg.py

 from __future__ import absolute_import, print_function, division
 import unittest
 
+import itertools
 import numpy as np
 import numpy.linalg
 from numpy.testing import assert_array_almost_equal
@@ -131,15 +132,42 @@ def test_qr_modes():
 def test_svd():
     rng = np.random.RandomState(utt.fetch_seed())
     A = tensor.matrix("A", dtype=theano.config.floatX)
-    U, V, T = svd(A)
-    fn = function([A], [U, V, T])
+    U, S, VT = svd(A)
+    fn = function([A], [U, S, VT])
     a = rng.rand(4, 4).astype(theano.config.floatX)
-    n_u, n_v, n_t = np.linalg.svd(a)
-    t_u, t_v, t_t = fn(a)
+    n_u, n_s, n_vt = np.linalg.svd(a)
+    t_u, t_s, t_vt = fn(a)
 
     assert _allclose(n_u, t_u)
-    assert _allclose(n_v, t_v)
-    assert _allclose(n_t, t_t)
+    assert _allclose(n_s, t_s)
+    assert _allclose(n_vt, t_vt)
+
+    fn = function([A], svd(A, compute_uv=False))
+    t_s = fn(a)
+    assert _allclose(n_s, t_s)
+
+
+def test_svd_infer_shape():
+    rng = np.random.RandomState(utt.fetch_seed())
+    A = tensor.matrix("A", dtype=theano.config.floatX)
+
+    for shp, full_matrices in itertools.product([(4, 4), (2, 4), (4, 2)],
+                                                [True, False]):
+        U, S, VT = svd(A, full_matrices=full_matrices)
+        fn = function([A], [U, S, VT])
+        fn_shp = function([A], [U.shape, S.shape, VT.shape])
+        a = rng.rand(*shp).astype(theano.config.floatX)
+        t_u, t_s, t_vt = fn(a)
+        t_u_shp, t_s_shp, t_vt_shp = fn_shp(a)
+
+        assert _allclose(t_u.shape, t_u_shp)
+        assert _allclose(t_s.shape, t_s_shp)
+        assert _allclose(t_vt.shape, t_vt_shp)
+
+    fn = function([A], svd(A, compute_uv=False))
+    fn_shp = function([A], svd(A, compute_uv=False).shape)
+    a = rng.rand(4, 2).astype(theano.config.floatX)
+    assert _allclose(fn(a).shape, fn_shp(a))
 
 
 def test_tensorsolve():