tensorinv op added, test for tensorinv added

theano/tensor/nlinalg.py

@@ -726,3 +726,59 @@ 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_dtype = a.dtype
+        out = theano.tensor.matrix(dtype=out_dtype)
+        return Apply(self, [a], [out])
+
+    def perform(self, node, inputs, outputs):
+        (a,) = inputs
+        (x,) = outputs
+        x[0] = self._numop(a, self.ind)
+
+
+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,8 @@ from theano.tensor.nlinalg import ( MatrixInverse,
                                     qr,
                                     matrix_power,
                                     norm,
-                                    svd
+                                    svd,
+                                    tensorinv
                                     )
 from nose.plugins.attrib import attr
 
@@ -517,3 +518,17 @@ 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)
+
+
+def test_tensorinv():
+    A = tensor.tensor4("A", dtype=theano.config.floatX)
+    X = tensorinv(A)
+    tf = function([A], [X])
+
+    a = numpy.eye(4 * 6)
+    a.shape = (4, 6, 8, 3)
+
+    n_ainv = numpy.linalg.tensorinv(a)
+    t_ainv = tf(a)
+
+    assert _allclose(n_ainv, t_ainv)