Repository: Theano/Theano

Issue:
Partial evaluation of a Theano function
This would allow to reuse the same function for test and train. With the cvm/vm linker, this is relatively easy if we compile the full function, then just ask to compute some of the outputs.

We could also allow to compute just some of the intermediate results. This would allow to unroll a big scan with the max supported lenght and compute just the number of iteration we want. Not sure how this would work for the back-ward.



Discussion:
For people wanting to know how to implement this:
- Understand how Theano graph are done. (bi-partite graph, with 2 type of node Variable and Apply node, you must understand how to navigate between each node)
- Start by the vm (python version). When this is done there, this can be ported to the cvm (c code). It is easier to work in python to develop the algorithm.
- start by understanding how the vm/cvm evaluate a Theano function. The VM is in the class VM, in the file theano/gof/vm.py. It start from the output, for each node, it check if its dependency are available, if not it try to evaluate this. This is done recursively, but with a loop implementation instead of calling recursively different function.
  -Then understand how the compute_map and storage_map work. In particular, understand how the IfElse op work, as this will cause only some of the inputs to be evaluated.
- With that information, you should be able to modify the VM to start from a subset of the output nodes. Implement a user way to do that. This would be the first user available feature. 

Could I ask for (even more) clarification on that? I mean, as partial evaluation you mean:
a) putting just a subset of inputs, calculate some subset of nodes and then, if you give the rest of the inputs multiple times, you don't recompute the nodes that were saved previously OR
b) give all the inputs, but calculate only some subset of the output nodes. This seems as more expected interpretation, as I see "output nodes" in the description, but how is that different from creating separate functions for different output nodes? Isn't it the case that you have to calculate the whole function exactly once for a given input anyway?
In particular, could you elaborate more on the "not recompiling train/test function"?
Thanks for the hints, sorry if these questions are silly.

edit: after a bit of thought, am I correct that b) is right and what you want to actually save is _just_ the compile time. The user can combine two functions into one and compile/optimize it once. Then, it applies the (potentially partial) input to it, to get partial output? So, e.g:

``` python
f = theano.function([a,b,c], [a**2*b, a**2 + c])
# compiling and optimizing once...
f(2, 3, ?)[0] == 12
f(1, ?, 3)[1] == 4
f(?, 1, 2)[0] # error
```

Don't hesitate to ask questions. It is only by well understanding the
problem that we can find a good solution:)

It is more b). Maybe you don' t need to give all inputs as some are maybe
not needed to compute the subset of output we need to compute.

This is similar to compiling many function, but would save on the
compilation time and memory consumption (not so much a problem). As the
goal is to allow many different combination of outputs and same compilation
time, we don't want to retrain a function each time.

On Thu, Mar 17, 2016 at 3:07 PM, Jakub Sygnowski notifications@github.com
wrote:

> Could I ask for (even more) clarification on that? I mean, as partial
> evaluation you mean:
> a) putting just a subset of inputs, calculate some subset of nodes and
> then, if you give the rest of the inputs multiple times, you don't
> recompute the nodes that were saved previously OR
> b) give all the inputs, but calculate only some subset of the output
> nodes. This seems as more expected interpretation, as I see "output nodes"
> in the description, but how is that different from creating separate
> functions for different output nodes? Isn't it the case that you have to
> calculate the whole function exactly once for a given input?
> In particular, could you elaborate more on the "not recompiling train/test
> function"?
> Thanks for the hints, sorry if these questions are silly.
> 
> —
> You are receiving this because you authored the thread.
> Reply to this email directly or view it on GitHub
> https://github.com/Theano/Theano/issues/2472#issuecomment-198037154


Resolving patch:
diff --git a/theano/compile/function_module.py b/theano/compile/function_module.py
--- a/theano/compile/function_module.py
+++ b/theano/compile/function_module.py
@@ -586,7 +586,8 @@ def copy(self, share_memory=False, swap=None, delete_updates=False,
 
         Returns
         -------
-        Copied theano.Function
+        theano.Function
+            Copied theano.Function
         """
         # helper function
         def checkSV(sv_ori, sv_rpl):
@@ -752,9 +753,37 @@ def checkSV(sv_ori, sv_rpl):
         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
+            The function inputs can be passed as keyword argument. For this, use
+            the name of the input or the input instance as the key.
+            Keyword argument ``output_subset`` is a list of either indices of the
+            function's outputs or the keys belonging to the `output_keys` dict
+            and represent outputs that are requested to be calculated.
+
+        Returns
+        -------
+        list
+            List of outputs on indices/keys from ``output_subset`` or all of them,
+            if ``output_subset`` is not passed.
+        """
         profile = self.profile
         t0 = time.time()
 
+        output_subset = kwargs.pop('output_subset', None)
+        if output_subset is not None and self.output_keys is not None:
+            output_subset =\
+                [self.output_keys.index(key) for key in output_subset]
+
         # Reinitialize each container's 'provided' counter
         if self.trust_input:
             i = 0
@@ -856,7 +885,9 @@ def __call__(self, *args, **kwargs):
         # Do the actual work
         t0_fn = time.time()
         try:
-            outputs = self.fn()
+            outputs =\
+                self.fn() if output_subset is None else\
+                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 +962,8 @@ def __call__(self, *args, **kwargs):
 
         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:
 
@@ -939,9 +971,16 @@ def __call__(self, *args, **kwargs):
 
                 assert len(self.output_keys) == len(outputs)
 
-                return dict(izip(self.output_keys, outputs))
+                if output_subset is None:
+                    return dict(izip(self.output_keys, outputs))
+                else:
+                    return dict((self.output_keys[index], outputs[index])
+                                for index in output_subset)
 
-            return outputs
+            if output_subset is None:
+                return outputs
+            else:
+                return [outputs[i] for i in output_subset]
 
     value = property(
         lambda self: self._value,
diff --git a/theano/gof/vm.py b/theano/gof/vm.py
--- a/theano/gof/vm.py
+++ b/theano/gof/vm.py
@@ -394,7 +394,7 @@ def run_thunk_of_node(self, node):
             )
         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 @@ def __call__(self):
             compute_map[k][0] = (k.owner is None)
 
         # apply_stack contains nodes
-        apply_stack = list(self.base_apply_stack)
+        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 @@ def __init__(self, allow_gc=None, use_cloop=False, callback=None,
         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 @@ def make_vm(self, nodes, thunks,
         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
+                getattr(self, 'allow_partial_eval', False)):
 
             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 getattr(self, 'allow_partial_eval', False):
+                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(


Test patch:
diff --git a/theano/gof/tests/test_vm.py b/theano/gof/tests/test_vm.py
--- a/theano/gof/tests/test_vm.py
+++ b/theano/gof/tests/test_vm.py
@@ -194,6 +194,28 @@ def time_linker(name, linker):
                                                        use_cloop=True))
 
 
+def test_partial_function():
+    import numpy as np
+    from theano.tests import unittest_tools as utt
+    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]]
+    utt.assert_allclose(f(5), np.array([32., 16., 1.7857142857142858]))
+
+
+def test_partial_function_output_keys():
+    x = tensor.scalar('input')
+    y = 3 * x
+    f = theano.function([x], {'a': y * 5, 'b': y - 7}, mode=Mode(
+        optimizer=None, linker=vm.VM_Linker(allow_partial_eval=True)))
+
+    assert f(5, output_subset=['a'])['a'] == f(5)['a']
+
+
 def test_allow_gc_cvm():
     mode = theano.config.mode
     if mode in ['DEBUG_MODE', 'DebugMode']: