Add tensorsolve to theano.tensor.nlinalg. Test.

theano/tensor/nlinalg.py

@@ -793,3 +793,75 @@ def tensorinv(a, ind=2):
         If `a` is singular or not 'square' (in the above sense).
     """
     return TensorInv(ind)(a)
+
+class TensorSolve(Op):
+    """
+    Theano utilization of numpy.linalg.tensorsolve
+
+    Solve the tensor equation ``a x = b`` for x.
+    It is assumed that all indices of `x` are summed over in the product,
+    together with the rightmost indices of `a`, as is done in, for example,
+    ``tensordot(a, x, axes=len(b.shape))``.
+    Parameters
+    ----------
+    a : array_like
+        Coefficient tensor, of shape ``b.shape + Q``. `Q`, a tuple, equals
+        the shape of that sub-tensor of `a` consisting of the appropriate
+        number of its rightmost indices, and must be such that
+        ``prod(Q) == prod(b.shape)`` (in which sense `a` is said to be
+        'square').
+    b : array_like
+        Right-hand tensor, which can be of any shape.
+    axes : tuple of ints, optional
+        Axes in `a` to reorder to the right, before inversion.
+        If None (default), no reordering is done.
+    """
+    _numop = staticmethod(numpy.linalg.tensorsolve)
+    __props__ = ('axes', )
+
+    def __init__(self, axes=None):
+        self.axes = axes
+
+    def make_node(self, a, b):
+        a = as_tensor_variable(a)
+        b = as_tensor_variable(b)
+        out_dtype = theano.scalar.upcast(a.dtype, b.dtype)
+        x = theano.tensor.matrix(dtype=out_dtype)
+        return Apply(self, [a, b], [x])
+
+    def perform(self, node, inputs, outputs):
+        (a, b,) = inputs
+        (x,) = outputs
+        x[0] = self._numop(a, b, self.axes)
+
+
+def tensorsolve(a, b, axes=None):
+    """
+    Theano utilization of numpy.linalg.tensorsolve
+
+    Solve the tensor equation ``a x = b`` for x.
+    It is assumed that all indices of `x` are summed over in the product,
+    together with the rightmost indices of `a`, as is done in, for example,
+    ``tensordot(a, x, axes=len(b.shape))``.
+    Parameters
+    ----------
+    a : array_like
+        Coefficient tensor, of shape ``b.shape + Q``. `Q`, a tuple, equals
+        the shape of that sub-tensor of `a` consisting of the appropriate
+        number of its rightmost indices, and must be such that
+        ``prod(Q) == prod(b.shape)`` (in which sense `a` is said to be
+        'square').
+    b : array_like
+        Right-hand tensor, which can be of any shape.
+    axes : tuple of ints, optional
+        Axes in `a` to reorder to the right, before inversion.
+        If None (default), no reordering is done.
+    Returns
+    -------
+    x : ndarray, shape Q
+    Raises
+    ------
+    LinAlgError
+        If `a` is singular or not 'square' (in the above sense).
+
+    return TensorSolve(axes)(a, b)

theano/tensor/tests/test_nlinalg.py

@@ -40,6 +40,7 @@ from theano.tensor.nlinalg import ( MatrixInverse,
                                     svd,
                                     TensorInv,
                                     tensorinv
+                                    tensorsolve
                                     )
 from nose.plugins.attrib import attr
 
@@ -162,6 +163,23 @@ def test_svd():
     assert _allclose(n_t, t_t)
 
 
+def test_tensorsolve():
+    rng = numpy.random.RandomState(utt.fetch_seed())
+    A = tensor.tensor4("A", dtype=theano.config.floatX)
+    B = tensor.matrix("B", dtype=theano.config.floatX)
+    X = tensorsolve(A, B)
+    fn = function([A, B], [X])
+
+    # slightly modified example from numpy.linalg.tensorsolve docstring
+    a = numpy.eye(2*3*4).astype(theano.config.floatX)
+    a.shape = (2*3, 4, 2, 3*4)
+    b = rng.rand(2*3, 4).astype(theano.config.floatX)
+
+    n_x = numpy.linalg.tensorsolve(a, b)
+    t_x = fn(a, b)
+    assert _allclose(n_x, t_x)
+
+
 def test_inverse_singular():
     singular = numpy.array([[1, 0, 0]] + [[0, 1, 0]] * 2,
                            dtype=theano.config.floatX)