partial function evaluation in python vm

theano/compile/function_module.py

@@ -752,9 +752,32 @@ class Function(object):
         return f_cpy
 
     def __call__(self, *args, **kwargs):
+        """
+        Evaluates value of a function on given arguments.
+
+        Parameters
+        ----------
+        args : list
+            List of inputs to the function. All inputs are required, even when
+            some of them are not necessary to calculate requested subset of
+            outputs.
+
+        kwargs : dict
+            TODO: other kwargs?
+            Keyword argument `output_subset` is a list of indices of the
+            function's outputs that are requested to be calculated.
+
+        Returns
+        -------
+        List of outputs on indices `output_subset` or all of them, if
+        `outpus_subset` is not passed. If there's only one output, returns just
+        the value.
+        """
         profile = self.profile
         t0 = time.time()
 
+        output_subset = kwargs.pop('output_subset', None)
+
         # Reinitialize each container's 'provided' counter
         if self.trust_input:
             i = 0
@@ -856,7 +879,7 @@ class Function(object):
         # Do the actual work
         t0_fn = time.time()
         try:
-            outputs = self.fn()
+            outputs = self.fn(output_subset=output_subset)
         except Exception:
             if hasattr(self.fn, 'position_of_error'):
                 # this is a new vm-provided function or c linker
@@ -931,7 +954,8 @@ class Function(object):
 
         if self.return_none:
             return None
-        elif self.unpack_single and len(outputs) == 1:
+        elif self.unpack_single and len(outputs) == 1 and\
+                output_subset is None:
             return outputs[0]
         else:
 
@@ -941,7 +965,10 @@ class Function(object):
 
                 return dict(izip(self.output_keys, outputs))
 
+            if output_subset is None:
                 return outputs
+            else:
+                return [outputs[i] for i in output_subset]
 
     value = property(
         lambda self: self._value,

theano/gof/tests/test_vm.py

@@ -194,6 +194,18 @@ def test_speed_lazy():
                                                        use_cloop=True))
 
 
+def test_partial_function():
+    import numpy as np
+    x = tensor.scalar('input')
+    y = x ** 2
+    f = theano.function([x], [y + 7, y - 9, y / 14.], mode=Mode(
+        optimizer=None, linker=vm.VM_Linker(allow_partial_eval=True)))
+
+    assert f(3, output_subset=[0, 1, 2]) == f(3)
+    assert f(4, output_subset=[0, 2]) == [f(4)[0], f(4)[2]]
+    assert abs(max(f(5) - np.array([32., 16., 1.7857142857142858]))) < 1e-9
+
+
 def test_allow_gc_cvm():
     mode = theano.config.mode
     if mode in ['DEBUG_MODE', 'DebugMode']:

theano/gof/vm.py

@@ -394,7 +394,7 @@ class Stack(VM):
             )
         return rval, dt
 
-    def __call__(self):
+    def __call__(self, output_subset=None):
         storage_map = self.storage_map
         compute_map = self.compute_map
         thunks = self.thunks
@@ -406,7 +406,13 @@ class Stack(VM):
             compute_map[k][0] = (k.owner is None)
 
         # apply_stack contains nodes
+        if output_subset is not None:
+            apply_stack =\
+                [self.outputs[i].owner for i in output_subset
+                    if self.outputs[i].owner]
+        else:
             apply_stack = list(self.base_apply_stack)
+
         last_apply_stack_len = -1
 
         # This record all function inputs/shared varibles and constants
@@ -680,11 +686,15 @@ class VM_Linker(link.LocalLinker):
     c_thunks
         If None or True, don't change the default. If False,
         don't compile c code for the thunks.
+    allow_partial_eval
+        If True, enforces usage of Stack or CVM, to allow for partial
+        evaluation of functions (calculating a subset of outputs).
 
     """
 
     def __init__(self, allow_gc=None, use_cloop=False, callback=None,
-                 lazy=None, schedule=None, c_thunks=None):
+                 lazy=None, schedule=None, c_thunks=None,
+                 allow_partial_eval=None):
         # Note: if more parameters are added to __init__, make sure to forward
         # them in the "type(self)(...)" call in the "accept" method below.
         if allow_gc is None:
@@ -695,6 +705,7 @@ class VM_Linker(link.LocalLinker):
         self.callback = callback
         self.lazy = lazy
         self.c_thunks = c_thunks
+        self.allow_partial_eval = allow_partial_eval
         self.updated_vars = {}
         if schedule:
             self.schedule = schedule
@@ -809,13 +820,18 @@ class VM_Linker(link.LocalLinker):
         pre_call_clear = [storage_map[v] for v in self.no_recycling]
 
         if (self.callback is not None or
-                (config.profile and config.profile_memory)):
+                (config.profile and config.profile_memory) or
+                self.allow_partial_eval):
 
             if self.use_cloop and self.callback is not None:
                 logger.warn('CVM does not support callback, using Stack VM.')
             if self.use_cloop and config.profile_memory:
                 warnings.warn(
                     'CVM does not support memory profile, using Stack VM.')
+            if self.use_cloop and self.allow_partial_eval:
+                warnings.warn(
+                    'CVM does not support partial evaluation yet, '
+                    'using Stack VM.')
             # Needed for allow_gc=True, profiling and storage_map reuse
             deps = self.compute_gc_dependencies(storage_map)
             vm = Stack(