Merge pull request #1335 from mrocklin/matrix_of_scalars

theano/tensor/basic.py

@@ -8224,3 +8224,33 @@ def diag(v, k=0):
         return diagonal(v, k)
     else:
         raise ValueError("Input must be 1- or 2-d.")
+
+def stacklists(arg):
+    """ Recursivly stack lists of tensors to maintain similar structure
+
+    This function can create a tensor from a shaped list of scalars
+    >>> from theano.tensor import stacklists, scalar, matrix
+    >>> from theano import function
+    >>> a,b,c,d = map(scalar, 'abcd')
+    >>> X = stacklists([[a, b],
+    ...                 [c, d]])
+    >>> f = function([a, b, c, d], X)
+    >>> f(1, 2, 3, 4)
+    array([[ 1.,  2.],
+           [ 3.,  4.]], dtype=float32)
+
+    We can also stack arbitrarily shaped tensors.  Here we stack matrices into
+    a 2 by 2 grid.
+    >>> from numpy import ones
+    >>> a,b,c,d, = [tensor.matrix(a) for a in 'abcd']
+    >>> X = stacklists([[a, b],
+    ...                 [c, d]])
+    >>> f = function([a, b, c, d], X)
+    >>> x = ones((4, 4), 'float32')
+    >>> f(x, x, x, x).shape
+    (2, 2, 4, 4)
+    """
+    if isinstance(arg, (tuple, list)):
+        return stack(*map(stacklists, arg))
+    else:
+        return arg

theano/tensor/tests/test_basic.py

@@ -43,7 +43,8 @@ from theano.tensor import (_shared, wvector, bvector, autocast_float_as,
         ScalarFromTensor, TensorFromScalar, dtensor4, Rebroadcast, Alloc,
         dtensor3, SpecifyShape, Mean, IncSubtensor, AdvancedIncSubtensor1,
         itensor3, Tile, AdvancedIncSubtensor, switch, Diagonal, Diag,
-        nonzero, flatnonzero, nonzero_values, inplace_increment)
+        nonzero, flatnonzero, nonzero_values, inplace_increment,
+        stacklists)
 from theano.tests import unittest_tools as utt
 
 
@@ -6778,6 +6779,32 @@ def test_transpose():
     assert tensor.transpose(x3).name == 'x3.T'
     assert tensor.transpose(tensor.dmatrix()).name is None
 
+def test_stacklists():
+    a,b,c,d = map(scalar, 'abcd')
+    X = stacklists([[a, b],
+                    [c, d]])
+    f = function([a, b, c, d], X)
+    result = f(1,2,3,4)
+    assert result.shape == (2, 2)
+    assert numpy.allclose(f(1, 2, 3, 4), numpy.asarray([[1,2],[3,4]]))
+
+    X = stacklists([a,b,c,d])
+    f = function([a, b, c, d], X)
+    result = f(1,2,3,4)
+    assert result.shape == (4,)
+    assert numpy.allclose(f(1, 2, 3, 4), numpy.asarray([[1,2,3,4]]))
+
+    X = stacklists([[[a],[b]],[[c],[d]]])
+    f = function([a, b, c, d], X)
+    result = f(1,2,3,4)
+    assert result.shape == (2, 2, 1)
+
+    a,b,c,d = [matrix(a) for a in 'abcd']
+    X = stacklists([[a, b],
+                    [c, d]])
+    f = function([a, b, c, d], X)
+    x = numpy.ones((4, 4), 'float32')
+    assert f(x,x,x,x).shape == (2, 2, 4, 4)
 
 class TestSpecifyShape(unittest.TestCase):
     def shortDescription(self):