Repository: Theano/Theano

Issue:
Partial function evaluation, implement in the cvm
Follow up of https://github.com/Theano/Theano/issues/2472. Make it work for the CVM, not just Stack.



Discussion:

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
@@ -735,9 +735,13 @@ def __call__(self, *args, **kwargs):
         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.
+            and represent outputs that are requested to be calculated. Regardless
+            of the presence of ``output_subset``, the updates are always calculated
+            and processed. To disable the updates, you should use the ``copy``
+            method with ``delete_updates=True``.
 
         Returns
         -------
diff --git a/theano/gof/lazylinker_c.py b/theano/gof/lazylinker_c.py
--- a/theano/gof/lazylinker_c.py
+++ b/theano/gof/lazylinker_c.py
@@ -15,7 +15,7 @@
 _logger = logging.getLogger('theano.gof.lazylinker_c')
 
 force_compile = False
-version = 0.21  # must match constant returned in function get_version()
+version = 0.211  # must match constant returned in function get_version()
 lazylinker_ext = None
 
 
diff --git a/theano/gof/vm.py b/theano/gof/vm.py
--- a/theano/gof/vm.py
+++ b/theano/gof/vm.py
@@ -332,7 +332,8 @@ def recursively_evaluate(var):
 
     def __init__(self, nodes, thunks, pre_call_clear,
                  storage_map, compute_map, fgraph, allow_gc,
-                 dependencies=None, callback=None, callback_input=None):
+                 n_updates, dependencies=None, callback=None,
+                 callback_input=None):
         super(Stack, self).__init__(nodes, thunks, pre_call_clear)
 
         self.allow_gc = allow_gc
@@ -346,6 +347,7 @@ def __init__(self, nodes, thunks, pre_call_clear,
         self.node_idx = node_idx = {}
         self.callback = callback
         self.callback_input = callback_input
+        self.n_updates = n_updates
 
         ords = fgraph.orderings()
 
@@ -417,6 +419,9 @@ def __call__(self, output_subset=None):
 
         # apply_stack contains nodes
         if output_subset is not None:
+            first_updated = len(self.outputs) - self.n_updates
+            output_subset = output_subset + list(range(first_updated,
+                                                       len(self.outputs)))
             apply_stack =\
                 [self.outputs[i].owner for i in output_subset
                     if self.outputs[i].owner]
@@ -425,7 +430,7 @@ def __call__(self, output_subset=None):
 
         last_apply_stack_len = -1
 
-        # This record all function inputs/shared varibles and constants
+        # This record all function inputs/shared variables and constants
         for var, data in iteritems(self.storage_map):
             if data[0] is None:
                 continue
@@ -842,7 +847,7 @@ def make_vm(self, nodes, thunks,
 
         if (self.callback is not None or self.callback_input is not None or
                 (config.profile and config.profile_memory) or
-                self.allow_partial_eval):
+                (self.allow_partial_eval and not self.use_cloop)):
 
             if self.use_cloop and (self.callback is not None or
                                    self.callback_input is not None):
@@ -850,9 +855,9 @@ def make_vm(self, nodes, thunks,
             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:
+            if not self.use_cloop and self.allow_partial_eval:
                 warnings.warn(
-                    'CVM does not support partial evaluation yet, '
+                    'LoopGC does not support partial evaluation, '
                     'using Stack VM.')
             # Needed for allow_gc=True, profiling and storage_map reuse
             deps = self.compute_gc_dependencies(storage_map)
@@ -860,6 +865,7 @@ def make_vm(self, nodes, thunks,
                 nodes, thunks, pre_call_clear,
                 storage_map, compute_map,
                 self.fgraph, self.allow_gc,
+                len(updated_vars),
                 dependencies=deps,
                 callback=self.callback,
                 callback_input=self.callback_input)
@@ -1000,7 +1006,8 @@ def make_vm(self, nodes, thunks,
                     nodes, thunks, pre_call_clear,
                     storage_map, compute_map,
                     self.fgraph, self.allow_gc,
-                    dependencies=deps
+                    len(updated_vars),
+                    dependencies=deps,
                 )
         return vm
 


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
@@ -197,23 +197,53 @@ def time_linker(name, linker):
 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 check_partial_function(linker_name):
+        x = tensor.scalar('input')
+        y = x ** 2
+        f = theano.function([x], [y + 7, y - 9, y / 14.], mode=Mode(
+            optimizer=None, linker=linker_name))
 
+        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)))
+    check_partial_function(vm.VM_Linker(allow_partial_eval=True, use_cloop=False))
+    check_partial_function('cvm')
 
-    assert f(5, output_subset=['a'])['a'] == f(5)['a']
+
+def test_partial_function_with_output_keys():
+
+    def check_partial_function_output_keys(linker_name):
+        x = tensor.scalar('input')
+        y = 3 * x
+        f = theano.function([x], {'a': y * 5, 'b': y - 7}, mode=Mode(
+            optimizer=None, linker=linker_name))
+
+        assert f(5, output_subset=['a'])['a'] == f(5)['a']
+
+    check_partial_function_output_keys(vm.VM_Linker(allow_partial_eval=True, use_cloop=False))
+    check_partial_function_output_keys('cvm')
+
+
+def test_partial_function_with_updates():
+
+    def check_updates(linker_name):
+        x = tensor.lscalar('input')
+        y = theano.shared(1, name='global')
+        f = theano.function([x], [x, x + 34], updates=[(y, x + 1)], mode=Mode(
+            optimizer=None, linker=linker_name))
+        g = theano.function([x], [x - 6], updates=[(y, y + 3)], mode=Mode(
+            optimizer=None, linker=linker_name))
+
+        assert f(3, output_subset=[]) == []
+        assert y.get_value() == 4
+        assert g(30, output_subset=[0]) == [24]
+        assert g(40, output_subset=[]) == []
+        assert y.get_value() == 10
+
+    check_updates(vm.VM_Linker(allow_partial_eval=True, use_cloop=False))
+    check_updates('cvm')
 
 
 def test_allow_gc_cvm():