Merge pull request #675 from mrocklin/shape-of-variables

theano/tensor/tests/test_utils.py

 import numpy
 
-from theano.tensor.utils import hash_from_ndarray, hash_from_dict
+import theano
+from theano.tensor.utils import (hash_from_ndarray, hash_from_dict,
+        shape_of_variables)
 
 
 def test_hash_from_ndarray():
@@ -46,3 +48,23 @@ def test_hash_from_dict():
 
     # List are not hashable. So they are transformed into tuple.
     assert hash_from_dict({0: (0,)}) == hash_from_dict({0: [0]})
+
+def test_shape_of_variables_simple():
+    x = theano.tensor.matrix('x')
+    y = x+x
+    env = theano.Env([x], [y])
+    assert shape_of_variables(env, {x: (5, 5)}) == {x: (5, 5), y: (5, 5)}
+
+    x = theano.tensor.matrix('x')
+    y = theano.tensor.dot(x, x.T)
+    env = theano.Env([x], [y])
+    shapes = shape_of_variables(env, {x: (5, 1)})
+    assert shapes[x] == (5, 1)
+    assert shapes[y] == (5, 5)
+
+def test_shape_of_variables_subtensor():
+    x = theano.tensor.matrix('x')
+    subx = x[1:]
+    env = theano.Env([x], [subx])
+    shapes = shape_of_variables(env, {x: (10, 10)})
+    assert shapes[subx] == (9, 10)

theano/tensor/utils.py

 import numpy
 
+import theano
 from theano.gof.cc import hash_from_code
 
 
@@ -43,3 +44,45 @@ def hash_from_dict(d):
             second_part += [v]
     tuple_items = tuple(first_part + second_part)
     return hash(tuple_items)
+
+def shape_of_variables(env, input_shapes):
+    """
+    Compute the numeric shape of all intermediate variables given input shapes
+
+    Inputs:
+        env - the theano.Env in question
+        input_shapes - a dict mapping input to shape
+
+    Outputs:
+        shapes - a dict mapping variable to shape
+
+    WARNING : This modifies the env. Not pure.
+
+    >>> import theano
+    >>> x = theano.tensor.matrix('x')
+    >>> y = x[512:]; y.name = 'y'
+    >>> env = theano.Env([x], [y])
+    >>> shape_of_variables(env, {x: (1024, 1024)})
+    {y: (512, 1024), x: (1024, 1024)}
+    """
+
+    if not hasattr(env, 'shape_feature'):
+        env.extend(theano.tensor.opt.ShapeFeature())
+
+    input_dims  = [dimension for inp in env.inputs
+                             for dimension in env.shape_feature.shape_of[inp]]
+
+    output_dims = [dimension for shape in env.shape_feature.shape_of.values()
+                             for dimension in shape]
+
+    compute_shapes = theano.function(input_dims, output_dims)
+
+    numeric_input_dims  = [dim for inp in env.inputs
+                               for dim in input_shapes[inp]]
+    numeric_output_dims = compute_shapes(*numeric_input_dims)
+
+    sym_to_num_dict = dict(zip(output_dims, numeric_output_dims))
+
+    return {var: tuple(sym_to_num_dict[sym]
+                             for sym in env.shape_feature.shape_of[var])
+                             for var in env.shape_feature.shape_of}