infer_shape for TensorInv added, test for it added

theano/tensor/nlinalg.py

@@ -742,7 +742,7 @@ class TensorInv(Op):
     def make_node(self, a):
         a = as_tensor_variable(a)
         out_dtype = a.dtype
-        out = theano.tensor.matrix(dtype=out_dtype)
+        out = theano.tensor.tensor4(dtype=out_dtype)
         return Apply(self, [a], [out])
 
     def perform(self, node, inputs, outputs):
@@ -750,6 +750,10 @@ class TensorInv(Op):
         (x,) = outputs
         x[0] = self._numop(a, self.ind)
 
+    def infer_shape(self, node, shapes):
+        sp = shapes[0][self.ind:] + shapes[0][:self.ind]
+        return [sp]
+
 
 def tensorinv(a, ind=2):
     """

theano/tensor/tests/test_nlinalg.py

@@ -38,6 +38,7 @@ from theano.tensor.nlinalg import ( MatrixInverse,
                                     matrix_power,
                                     norm,
                                     svd,
+                                    TensorInv,
                                     tensorinv
                                     )
 from nose.plugins.attrib import attr
@@ -520,15 +521,24 @@ class T_NormTests(unittest.TestCase):
             assert _allclose(n_n, t_n)
 
 
-def test_tensorinv():
-    A = tensor.tensor4("A", dtype=theano.config.floatX)
-    X = tensorinv(A)
-    tf = function([A], [X])
-
-    a = numpy.eye(4 * 6).astype(theano.config.floatX)
-    a.shape = (4, 6, 8, 3)
+class test_TensorInv(utt.InferShapeTester):
+    def setUp(self):
+        super(test_TensorInv, self).setUp()
+        self.A = tensor.tensor4("A", dtype=theano.config.floatX)
+        self.a = numpy.random.rand(4, 6, 8, 3).astype(theano.config.floatX)
 
-    n_ainv = numpy.linalg.tensorinv(a)
-    t_ainv = tf(a)
+    def test_infer_shape(self):
+        A = self.A
+        Ai = tensorinv(A)
+        self._compile_and_check([A],  # theano.function inputs
+                                [Ai],  # theano.function outputs
+                                [self.a],  # value to substitute
+                                TensorInv)
 
-    assert _allclose(n_ainv, t_ainv)
+    def test_eval(self):
+        A = self.A
+        Ai = tensorinv(A)
+        n_ainv = numpy.linalg.tensorinv(self.a)
+        tf = function([A], [Ai])
+        t_ainv = tf(self.a)
+        assert _allclose(n_ainv, t_ainv)