generalize tensor_of_scalars to stacklists

theano/tensor/basic.py

@@ -8225,20 +8225,32 @@ def diag(v, k=0):
     else:
         raise ValueError("Input must be 1- or 2-d.")
 
-def tensor_of_scalars(arg):
-    """ Returns a tensor of scalars
+def stacklists(arg):
+    """ Recursivly stack lists of tensors to maintain similar structure
 
-    >>> from theano.tensor import tensor_of_scalars, scalar
+    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 = tensor_of_scalars([[a, b],
-    ...                        [c, d]])
+    >>> 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(tensor_of_scalars, arg))
+        return stack(*map(stacklists, arg))
     else:
         return arg

theano/tensor/tests/test_basic.py

@@ -44,7 +44,7 @@ from theano.tensor import (_shared, wvector, bvector, autocast_float_as,
         dtensor3, SpecifyShape, Mean, IncSubtensor, AdvancedIncSubtensor1,
         itensor3, Tile, AdvancedIncSubtensor, switch, Diagonal, Diag,
         nonzero, flatnonzero, nonzero_values, inplace_increment,
-        tensor_of_scalars)
+        stacklists)
 from theano.tests import unittest_tools as utt
 
 
@@ -6779,26 +6779,33 @@ def test_transpose():
     assert tensor.transpose(x3).name == 'x3.T'
     assert tensor.transpose(tensor.dmatrix()).name is None
 
-def test_tensor_of_scalars():
+def test_stacklists():
     a,b,c,d = map(scalar, 'abcd')
-    X = tensor_of_scalars([[a, b],
-                           [c, d]])
+    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 = tensor_of_scalars([a,b,c,d])
+    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 = tensor_of_scalars([[[a],[b]],[[c],[d]]])
+    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):
         return None