Merge pull request #4546 from iikulikov/master

theano/tensor/nlinalg.py

@@ -726,3 +726,62 @@ def norm(x, ord):
             raise ValueError(0)
     elif ndim > 2:
         raise NotImplementedError("We don't support norm witn ndim > 2")
+
+
+class TensorInv(Op):
+    """
+    Class wrapper for tensorinv() function;
+    Theano utilization of numpy.linalg.tensorinv;
+    """
+    _numop = staticmethod(numpy.linalg.tensorinv)
+    __props__ = ('ind',)
+
+    def __init__(self, ind=2):
+        self.ind = ind
+
+    def make_node(self, a):
+        a = as_tensor_variable(a)
+        out = a.type()
+        return Apply(self, [a], [out])
+
+    def perform(self, node, inputs, outputs):
+        (a,) = inputs
+        (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):
+    """
+    Does not run on GPU;
+    Theano utilization of numpy.linalg.tensorinv;
+
+    Compute the 'inverse' of an N-dimensional array.
+    The result is an inverse for `a` relative to the tensordot operation
+    ``tensordot(a, b, ind)``, i. e., up to floating-point accuracy,
+    ``tensordot(tensorinv(a), a, ind)`` is the "identity" tensor for the
+    tensordot operation.
+
+    Parameters
+    ----------
+    a : array_like
+        Tensor to 'invert'. Its shape must be 'square', i. e.,
+        ``prod(a.shape[:ind]) == prod(a.shape[ind:])``.
+    ind : int, optional
+        Number of first indices that are involved in the inverse sum.
+        Must be a positive integer, default is 2.
+
+    Returns
+    -------
+    b : ndarray
+        `a`'s tensordot inverse, shape ``a.shape[ind:] + a.shape[:ind]``.
+
+    Raises
+    ------
+    LinAlgError
+        If `a` is singular or not 'square' (in the above sense).
+    """
+    return TensorInv(ind)(a)

theano/tensor/tests/test_nlinalg.py

@@ -37,7 +37,9 @@ from theano.tensor.nlinalg import ( MatrixInverse,
                                     qr,
                                     matrix_power,
                                     norm,
-                                    svd
+                                    svd,
+                                    TensorInv,
+                                    tensorinv
                                     )
 from nose.plugins.attrib import attr
 
@@ -517,3 +519,41 @@ class T_NormTests(unittest.TestCase):
             t_n = f(A[2][i])
             n_n = numpy.linalg.norm(A[2][i], A[3][i])
             assert _allclose(n_n, t_n)
+
+
+class test_TensorInv(utt.InferShapeTester):
+    def setUp(self):
+        super(test_TensorInv, self).setUp()
+        self.A = tensor.tensor4("A", dtype=theano.config.floatX)
+        self.B = tensor.tensor3("B", dtype=theano.config.floatX)
+        self.a = numpy.random.rand(4, 6, 8, 3).astype(theano.config.floatX)
+        self.b = numpy.random.rand(2, 15, 30).astype(theano.config.floatX)
+        self.b1 = numpy.random.rand(30, 2, 15).astype(theano.config.floatX)  # for ind=1 since we need prod(b1.shape[:ind]) == prod(b1.shape[ind:])
+
+    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)
+
+    def test_eval(self):
+        A = self.A
+        Ai = tensorinv(A)
+        n_ainv = numpy.linalg.tensorinv(self.a)
+        tf_a = function([A], [Ai])
+        t_ainv = tf_a(self.a)
+        assert _allclose(n_ainv, t_ainv)
+
+        B = self.B
+        Bi = tensorinv(B)
+        Bi1 = tensorinv(B, ind=1)
+        n_binv = numpy.linalg.tensorinv(self.b)
+        n_binv1 = numpy.linalg.tensorinv(self.b1, ind=1)
+        tf_b = function([B], [Bi])
+        tf_b1 = function([B], [Bi1])
+        t_binv = tf_b(self.b)
+        t_binv1 = tf_b1(self.b1)
+        assert _allclose(t_binv, n_binv)
+        assert _allclose(t_binv1, n_binv1)