Merge pull request #4529 from yobibyte/tensorsolve

theano/tensor/nlinalg.py

@@ -793,3 +793,62 @@ 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
+    Class wrapper for tensorsolve function.
+
+    """
+    _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. Does not run on GPU!
+
+    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

@@ -39,7 +39,8 @@ from theano.tensor.nlinalg import ( MatrixInverse,
                                     norm,
                                     svd,
                                     TensorInv,
-                                    tensorinv
+                                    tensorinv,
+                                    tensorsolve
                                     )
 from nose.plugins.attrib import attr
 
@@ -162,6 +163,52 @@ 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)
+
+    # check the type upcast now
+    C = tensor.tensor4("C", dtype='float32')
+    D = tensor.matrix("D", dtype='float64')
+    Y = tensorsolve(C, D)
+    fn = function([C, D], [Y])
+
+    c = numpy.eye(2 * 3 * 4, dtype='float32')
+    c.shape = (2 * 3, 4, 2, 3 * 4)
+    d = rng.rand(2 * 3, 4).astype('float64')
+    n_y = numpy.linalg.tensorsolve(c, d)
+    t_y = fn(c, d)
+    assert _allclose(n_y, t_y)
+    assert n_y.dtype == Y.dtype
+
+    # check the type upcast now
+    E = tensor.tensor4("E", dtype='int32')
+    F = tensor.matrix("F", dtype='float64')
+    Z = tensorsolve(E, F)
+    fn = function([E, F], [Z])
+
+    e = numpy.eye(2 * 3 * 4, dtype='int32')
+    e.shape = (2 * 3, 4, 2, 3 * 4)
+    f = rng.rand(2 * 3, 4).astype('float64')
+    n_z = numpy.linalg.tensorsolve(e, f)
+    t_z = fn(e, f)
+    assert _allclose(n_z, t_z)
+    assert n_z.dtype == Z.dtype
+
+
 def test_inverse_singular():
     singular = numpy.array([[1, 0, 0]] + [[0, 1, 0]] * 2,
                            dtype=theano.config.floatX)