Merge pull request #212 from dwf/rebased_new_lazy_ifelse

doc/cifarSC2011/advanced_theano.txt

@@ -19,7 +19,7 @@ Conditions
 .. code-block:: python
 
   from theano import tensor as T
-  from theano.lazycond import ifelse
+  from theano.ifelse import ifelse
   import theano, time, numpy
 
   a,b = T.scalars('a','b')

doc/tutorial/conditions.txt

@@ -18,7 +18,7 @@ Conditions
 .. code-block:: python
 
   from theano import tensor as T
-  from theano.lazycond import ifelse
+  from theano.ifelse import ifelse
   import theano, time, numpy
 
   a,b = T.scalars('a','b')

theano/gof/tests/test_lazy.py

@@ -94,7 +94,7 @@ class NotImplementedOp(PureOp):
 
 
 def test_ifelse():
-    a = generic()
+    a = T.scalar()
     b = generic()
     c = generic()
 
@@ -105,15 +105,15 @@ def test_ifelse():
 
     try:
         print "case 1"
-        f( True, 'a', 'b')
+        f( 1, 'a', 'b')
         assert False
     except NotImplementedOp.E:
         pass
     print "... passed"
 
     print "case 2"
-    print f( False, 'a', 'b')
-    assert f( False, 'a', 'b') == 'b'
+    print f( 0, 'a', 'b')
+    assert f( 0, 'a', 'b') == 'b'
     print "... passed"
 
 
@@ -123,8 +123,8 @@ def more_complex_test():
 
     x1 = T.scalar('x1')
     x2 = T.scalar('x2')
-    c1 = generic('c1')
-    c2 = generic('c2')
+    c1 = T.scalar('c1')
+    c2 = T.scalar('c2')
     t1 = ifelse(c1,x1,notimpl(x2))
     t1.name = 't1'
     t2 = t1*10

theano/ifelse.py

+"""
+IfElse introduces lazy evaluation in Theano (coupled with the CVM/VM
+linkers). It resembles the if clause of any programming language, that
+has a `then` and `else` branch, and executes either one or the other
+according to the condition provided.
+
+This op contrast the already existent `switch` op, that will evaluate both
+branches of the clause and afterwards pick (according to the condition)
+which value to report. Note also that `switch` is an elemwise operation (so
+it picks each entry of a matrix according to the condition) while `ifelse`
+is a global operation with a scalar condition.
+"""
+
+__docformat__ = 'restructedtext en'
+__authors__ = ("Razvan Pascanu "
+               "James Bergstra "
+               "Dumitru Erhan "
+               "David Warde-Farley")
+__copyright__ = "(c) 2010, Universite de Montreal"
+__contact__ = "Razvan Pascanu <r.pascanu@gmail>"
+
+from copy import deepcopy
+from itertools import izip
+import logging
+
+from theano.gof import PureOp, Apply
+
+import theano.tensor
+import gof
+
+from compile import optdb
+from tensor import opt
+from scan_module.scan_utils import find_up
+from scan_module.scan_utils import clone
+
+_logger = logging.getLogger('theano.ifelse')
+
+
+class IfElse(PureOp):
+    """
+    Op that provides conditional graph evaluation if used with the CVM/VM
+    linkers. Note that there exist a helpful function `ifelse` that should
+    be used to instantiate the op!
+
+    According to a scalar condition `condition` the op evaluates and then
+    returns all the tensors provided on the `then` branch, otherwise it
+    evaluates and returns the tensors provided on the `else` branch. The op
+    supports multiple tensors on each branch, with the condition that the same
+    number of tensors are on the `then` as on the `else` and there is a one
+    to one correspondence between them (shape and dtype wise).
+
+    The `then` branch is defined as the first N tensors (after the
+    condition), while the `else` branch is defined as the last N tensors.
+
+    Example usage:
+
+        ``rval = ifelse(condition, rval_if_true1, .., rval_if_trueN,
+                        rval_if_false1, rval_if_false2, .., rval_if_falseN)``
+
+    :note:
+        Other Linkers then CVM and VM are INCOMPATIBLE with this Op, and
+        will ingnore its lazy characteristic, computing both the True and
+        False branch before picking one.
+
+    """
+    def __init__(self, n_outs, as_view=False, gpu=False, name=None):
+        if as_view:
+            # check destroyhandler and others to ensure that a view_map with
+            # multiple inputs can work
+            view_map = {}
+            for idx in xrange(n_outs):
+                view_map[idx] = [idx + 1]
+            self.view_map = view_map
+        self.as_view = as_view
+        self.gpu = gpu
+        self.n_outs = n_outs
+        self.name = name
+
+    def __eq__(self, other):
+        if not type(self) == type(other):
+            return False
+        if not self.as_view == other.as_view:
+            return False
+        if not self.gpu == other.gpu:
+            return False
+        if not self.n_outs == other.n_outs:
+            return False
+        return True
+
+    def __hash__(self):
+        rval = (hash(type(self)) ^
+                hash(self.as_view) ^
+                hash(self.gpu) ^
+                hash(self.n_outs))
+        return rval
+
+    def __str__(self):
+        args = []
+        if self.name is not None:
+            args.append(self.name)
+        if self.as_view:
+            args.append('inplace')
+        if self.gpu:
+            args.append('gpu')
+        return 'if{%s}' % ','.join(args)
+
+    def infer_shape(self, node, inputs_shapes):
+        # By construction, corresponding then/else pairs have the same number
+        # of dimensions
+
+        ts_shapes = inputs_shapes[1:][:self.n_outs]
+        fs_shapes = inputs_shapes[1:][self.n_outs:]
+        # All elements of all shape tuples for the true and false outputs are
+        # unpacked into the inputs of a separate ifelse, and then the outputs
+        # of that ifelse are packed back into shape tuples.
+        new_ts_inputs = []
+        for ts_shape in ts_shapes:
+            if isinstance(ts_shape, (list, tuple)):
+                new_ts_inputs += list(ts_shape)
+            else:
+                # It can be None for generic objects
+                return [None] * self.n_outs
+
+        new_fs_inputs = []
+        for fs_shape in fs_shapes:
+            if isinstance(fs_shape, (list, tuple)):
+                new_fs_inputs += list(fs_shape)
+            else:
+                # It can be None for generic objects
+                return [None] * self.n_outs
+
+        assert len(new_ts_inputs) == len(new_fs_inputs)
+        if len(new_ts_inputs + new_fs_inputs) > 0:
+            name_tokens = ['shape']
+            if self.name is not None:
+                name_tokens.append(self.name)
+
+            new_ifelse = IfElse(
+                n_outs=len(new_ts_inputs),
+                as_view=False,
+                gpu=False,
+                name='_'.join(name_tokens))
+            new_outs = new_ifelse.make_node(node.inputs[0],
+                            *(new_ts_inputs + new_fs_inputs)).outputs
+        else:
+            new_outs = []
+
+        # generate pairs of shapes
+        out_shapes = []
+        for out in node.outputs:
+            out_shapes.append(tuple(new_outs[:out.ndim]))
+            new_outs = new_outs[out.ndim:]
+
+        # new_outs should be an empty list after last iteration
+        assert len(new_outs) == 0
+
+        return out_shapes
+
+    def make_node(self, c, *args):
+        assert len(args) == 2 * self.n_outs, (
+                "Wrong number of arguments to make_node: "
+                "expected %d, got %d" % (2 * self.n_outs, len(args))
+        )
+        if not self.gpu:
+            # When gpu is true, we are given only cuda ndarrays, and we want
+            # to keep them be cuda ndarrays
+            c = theano.tensor.as_tensor_variable(c)
+            nw_args = []
+            for x in args:
+                if isinstance(x, theano.Variable):
+                    nw_args.append(x)
+                else:
+                    nw_args.append(theano.tensor.as_tensor_variable(x))
+            args = nw_args
+        ts = args[:self.n_outs]
+        fs = args[self.n_outs:]
+
+        for t, f in izip(ts, fs):
+            if t.type != f.type:
+                raise TypeError(('IfElse requires same types for true and '
+                                'false return values'), t, f, t.type, f.type)
+        if c.ndim > 0:
+            raise TypeError(('Condition given to the op has to be a scalar '
+                             'with 0 standing for False, anything else '
+                             'for True'))
+        return Apply(self, [c] + list(args), [t.type() for t in ts])
+
+    def R_op(self, inputs, eval_points):
+        return self.make_node(inputs[0], *eval_points[1:]).outputs
+
+    def grad(self, ins, grads):
+        ts = ins[1:][:self.n_outs]
+        fs = ins[1:][self.n_outs:]
+        if self.name is not None:
+            nw_name_t = self.name + '_grad_t'
+            nw_name_f = self.name + '_grad_f'
+        else:
+            nw_name_t = None
+            nw_name_f = None
+        if_true_op = IfElse(n_outs=self.n_outs,
+                            as_view=self.as_view,
+                            gpu=self.gpu,
+                            name=nw_name_t)
+
+        if_false_op = IfElse(n_outs=self.n_outs,
+                             as_view=self.as_view,
+                             gpu=self.gpu,
+                             name=nw_name_f)
+
+        if_true = ([ins[0]] + grads + [theano.tensor.zeros_like(t)
+                                       for t in ts])
+        if_false = ([ins[0]] + [theano.tensor.zeros_like(f)
+                                for f in fs] + grads)
+        return ([None] +
+                if_true_op.make_node(*if_true).outputs +
+                if_false_op.make_node(*if_false).outputs)
+
+    def make_thunk(self, node, storage_map, compute_map, no_recycling):
+        outtypes = [out.type for out in node.outputs]
+        cond = node.inputs[0]
+        ts = node.inputs[1:][:self.n_outs]
+        fs = node.inputs[1:][self.n_outs:]
+        outputs = node.outputs
+
+        def thunk():
+            if not compute_map[cond][0]:
+                return [0]
+            else:
+                truthval = storage_map[cond][0]
+                if truthval != 0:
+                    ls = [idx + 1 for idx in xrange(self.n_outs)
+                          if not compute_map[ts[idx]][0]]
+                    if len(ls) > 0:
+                        return ls
+                    else:
+                        for out, outtype, t in izip(outputs, outtypes, ts):
+                            compute_map[out][0] = 1
+                            if self.as_view:
+                                oval = outtype.filter(storage_map[t][0])
+                            else:
+                                oval = outtype.filter(
+                                    deepcopy(storage_map[t][0]))
+                            storage_map[out][0] = oval
+                        return []
+                else:
+                    ls = [1 + idx + self.n_outs for idx in xrange(self.n_outs)
+                          if not compute_map[fs[idx]][0]]
+                    if len(ls) > 0:
+                        return ls
+                    else:
+                        for out, outtype, f in izip(outputs, outtypes, fs):
+                            compute_map[out][0] = 1
+                            # can't view both outputs unless destroyhandler
+                            # improves
+                            oval = outtype.filter(
+                                deepcopy(storage_map[f][0]))
+                            storage_map[out][0] = oval
+                        return []
+
+        thunk.lazy = True
+        thunk.inputs = [storage_map[v] for v in node.inputs]
+        thunk.outputs = [storage_map[v] for v in node.outputs]
+        return thunk
+
+
+def ifelse(condition, then_branch, else_branch, name=None):
+    """
+    This function corresponds to an if statement, returning (and evaluating)
+    inputs in the ``then_branch`` if ``condition`` evaluates to True or
+    inputs in the ``else_branch`` if ``condition`` evalutates to False.
+
+    :type condition: scalar like
+    :param condition:
+        ``condition`` should be a tensor scalar representing the condition.
+        If it evaluates to 0 it corresponds to False, anything else stands
+        for True.
+
+    :type then_branch: list of theano expressions/ theano expressions
+    :param then_branch:
+        A single theano variable or a list of theano variables that the
+        function should return as the output if ``condition`` evaluates to
+        true. The number of variables should match those in the
+        ``else_branch``, and there should be a one to one correspondance
+        (type wise) with the tensors provided in the else branch
+
+    :type else_branch: list of theano expressions/ theano expressions
+    :param else_branch:
+        A single theano variable or a list of theano variables that the
+        function should return as the output if ``condition`` evaluates to
+        false. The number of variables should match those in the then branch,
+        and there should be a one to one correspondace (type wise) with the
+        tensors provided in the then branch.
+
+    :return:
+        A list of theano variables or a single variable (depending on the
+        nature of the ``then_branch`` and ``else_branch``). More exactly if
+        ``then_branch`` and ``else_branch`` is a tensor, then
+        the return variable will be just a single variable, otherwise a
+        list. The value returns correspond either to the values in the
+        ``then_branch`` or in the ``else_branch`` depending on the value of
+        ``cond``.
+    """
+    if type(then_branch) is not type(else_branch):
+        raise ValueError(('The two branches should be identical. '
+                          'This error could be raised if for example '
+                          ' you provided a one element list on the then '
+                          ' branch but a tensor on the else branch'))
+
+    rval_type = None
+    if type(then_branch) is list:
+        rval_type = list
+    elif type(then_branch) is tuple:
+        rval_type = tuple
+
+    if type(then_branch) not in (list, tuple):
+        then_branch = [then_branch]
+    if type(else_branch) not in (list, tuple):
+        else_branch = [else_branch]
+
+    if len(then_branch) != len(else_branch):
+        raise ValueError(('The number of values on the `then` branch'
+                          ' should have the same number of variables as '
+                          'the `else` branch : (variables on `then` '
+                          '%d' % len(then_branch) + ', variables on `else` '
+                          '%d' % len(else_branch) + ')'))
+
+    new_ifelse = IfElse(n_outs=len(then_branch),
+                        as_view=False,
+                        gpu=False,
+                        name=name)
+
+    ins = [condition] + list(then_branch) + list(else_branch)
+    rval = new_ifelse.make_node(*ins).outputs
+
+    if rval_type is None:
+        return rval[0]
+    elif rval_type is list:
+        return list(rval)
+    else:
+        return tuple(rval)
+
+
+@gof.local_optimizer([None])
+def cond_make_inplace(node):
+    op = node.op
+    if isinstance(op, IfElse) and not op.as_view:
+        return IfElse(n_outs=op.n_outs,
+                      as_view=True,
+                      gpu=op.gpu,
+                      name=op.name).make_node(*node.inputs).outputs
+    return False
+
+
+# XXX: Optimizations commented pending further debugging (certain optimizations
+# make computation less lazy than it should be currently).
+# optdb.register('cond_make_inplace', opt.in2out(cond_make_inplace,
+#     ignore_newtrees=True), 95, 'fast_run', 'inplace')
+#
+# ifelse_equilibrium = gof.EquilibriumDB()
+# ifelse_seqopt = gof.SequenceDB()
+# ifelse_equilibrium.register('seq_ifelse', ifelse_seqopt, 'fast_run',
+#                             'ifelse')
+''' Comments:
+I've wrote this comments to explain how the optimization of ifelse function
+(for future developers that need to parse this part of code. Please try to
+keep this comments in sync with whatever changes you add to the code.
+
+ifelse optimization are registered before canonicalize !
+
+The optimizations are called in sequence as follows:
+    * equilibrium shell (runs until no change):
+        * ifelse_lift
+        * ifelse_merge_ifs
+        * ifelse_merge_nodes
+        * ifelse_remove_identical_inside
+        * ifelse_sameCondTrue_inside
+        * ifelse_sameCondFalse_inside
+    * merge_nodes_1
+    * ifelse_sameCondTrue
+    * ifelse_sameCondFalse
+    * ifelse_removeIdentical
+
+where, each of the optimization do the following things:
+    `ifelse_lift` (def cond_lift_single_if):
+
+'''
+# optdb.register('ifelse_equilibriumOpt', ifelse_equilibrium, .5, 'fast_run',
+#                'ifelse')
+
+acceptable_ops = (theano.tensor.basic.Dot,
+                  theano.tensor.basic.Reshape,
+                  theano.tensor.basic.Shape,
+                  theano.tensor.basic.SpecifyShape,
+                  theano.tensor.basic.MaxAndArgmax,
+                  theano.tensor.basic.Subtensor,
+                  theano.tensor.basic.IncSubtensor,
+                  theano.tensor.basic.Rebroadcast,
+                  theano.tensor.basic.Alloc,
+                  theano.tensor.elemwise.Elemwise,
+                  theano.tensor.elemwise.DimShuffle)
+
+
+@gof.local_optimizer([None])
+def ifelse_lift_single_if_through_acceptable_ops(main_node):
+    """This optimization lifts up certain ifelse instances.
+
+        op(ifelse(c, x, y)) -> ifelse(c, op(x), op(y))
+
+    if `op` is in the `acceptable_ops` list, and there is no other if as
+    input to that specific `op`, and the if has no other clients !?
+    """
+    if not (isinstance(main_node.op, acceptable_ops)):
+        return False
+    all_inp_nodes = set()
+    for inp in main_node.inputs:
+        all_inp_nodes.add(inp.owner)
+    ifnodes = [x for x in list(all_inp_nodes)
+                    if x and isinstance(x.op, IfElse)]
+    # if we have multiple ifs as inputs .. it all becomes quite complicated
+    # :)
+    if len(ifnodes) != 1:
+        return False
+    node = ifnodes[0]
+    op = node.op
+
+    ts = node.inputs[1:][:op.n_outs]
+    fs = node.inputs[1:][op.n_outs:]
+
+    outs = main_node.outputs
+    mop = main_node.op
+    true_ins = []
+    false_ins = []
+
+    for x in main_node.inputs:
+        if x in node.outputs:
+            idx = node.outputs.index(x)
+            true_ins.append(ts[idx])
+            false_ins.append(fs[idx])
+        else:
+            true_ins.append(x)
+            false_ins.append(x)
+    true_eval = mop.make_node(*true_ins).outputs
+    false_eval = mop.make_node(*false_ins).outputs
+    #true_eval  = clone(outs, replace = dict(zip(node.outputs, ts)))
+    #false_eval = clone(outs, replace = dict(zip(node.outputs, fs)))
+
+    nw_outs = ifelse(node.inputs[0], true_eval, false_eval)
+    if type(nw_outs) not in (tuple, list):
+        nw_outs = [nw_outs]
+    return nw_outs
+
+
+@gof.local_optimizer([None])
+def cond_merge_ifs_true(node):
+    op = node.op
+    if not isinstance(op, IfElse):
+        return False
+    t_ins = node.inputs[1:][:op.n_outs]
+
+    replace = {}
+    for idx, tval in enumerate(t_ins):
+        if (tval.owner and isinstance(tval.owner.op, IfElse) and
+            tval.owner.inputs[0] == node.inputs[0]):
+            ins_op = tval.owner.op
+            ins_t = tval.owner.inputs[1:][:ins_op.n_outs]
+            replace[idx + 1] = ins_t[tval.owner.outputs.index(tval)]
+
+    if len(replace.items()) == 0:
+        return False
+
+    old_ins = list(node.inputs)
+    for pos, var in replace.items():
+        old_ins[pos] = var
+    return op.make_node(*old_ins).outputs
+
+
+@gof.local_optimizer([None])
+def cond_merge_ifs_false(node):
+    op = node.op
+    if not isinstance(op, IfElse):
+        return False
+    f_ins = node.inputs[1:][op.n_outs:]
+
+    replace = {}
+    for idx, fval in enumerate(f_ins):
+        if (fval.owner and isinstance(fval.owner.op, IfElse) and
+            fval.owner.inputs[0] == node.inputs[0]):
+            ins_op = fval.owner.op
+            ins_t = fval.owner.inputs[1:][ins_op.n_outs:]
+            replace[idx + 1 + op.n_outs] = \
+                    ins_t[fval.owner.outputs.index(fval)]
+
+    if len(replace.items()) == 0:
+        return False
+
+    old_ins = list(node.inputs)
+    for pos, var in replace.items():
+        old_ins[pos] = var
+    return op.make_node(*old_ins).outputs
+
+
+class CondMerge(gof.Optimizer):
+    """ Graph Optimizer that merges different cond ops """
+    def add_requirements(self, env):
+        env.extend(gof.toolbox.ReplaceValidate())
+
+    def apply(self, env):
+        nodelist = list(env.toposort())
+        cond_nodes = filter(lambda s: isinstance(s.op, IfElse), nodelist)
+        if len(cond_nodes) < 2:
+            return False
+        merging_node = cond_nodes[0]
+        for proposal in cond_nodes[1:]:
+            if (proposal.inputs[0] == merging_node.inputs[0] and
+                not find_up(proposal, merging_node)):
+                # Create a list of replacements for proposal
+                mn_ts = merging_node.inputs[1:][:merging_node.op.n_outs]
+                mn_fs = merging_node.inputs[1:][merging_node.op.n_outs:]
+                pl_ts = proposal.inputs[1:][:proposal.op.n_outs]
+                pl_fs = proposal.inputs[1:][proposal.op.n_outs:]
+                new_ins = ([merging_node.inputs[0]] +
+                           mn_ts + pl_ts + mn_fs + pl_fs)
+                mn_name = '?'
+                if merging_node.op.name:
+                    mn_name = merging_node.op.name
+                pl_name = '?'
+                mn_n_ts = len(mn_ts)
+                mn_n_fs = len(mn_fs)
+                if proposal.op.name:
+                    pl_name = proposal.op.name
+                new_ifelse = IfElse(
+                    n_outs=len(mn_ts + pl_ts),
+                    as_view=False,
+                    gpu=False,
+                    name=mn_name + '&' + pl_name)
+                print 'here'
+                new_outs = new_ifelse.make_node(*new_ins).outputs
+                new_outs = [clone(x) for x in new_outs]
+                old_outs = []
+                if type(merging_node.outputs) not in (list, tuple):
+                    old_outs += [merging_node.outputs]
+                else:
+                    old_outs += merging_node.outputs
+                if type(proposal.outputs) not in (list, tuple):
+                    old_outs += [proposal.outputs]
+                else:
+                    old_outs += proposal.outputs
+                pairs = zip(old_outs, new_outs)
+                env.replace_all_validate(pairs, reason='cond_merge')
+
+
+@gof.local_optimizer([None])
+def cond_remove_identical(node):
+    op = node.op
+
+    if not isinstance(op, IfElse):
+        return False
+    ts = node.inputs[1:][:op.n_outs]
+    fs = node.inputs[1:][op.n_outs:]
+
+    # sync outs
+    out_map = {}
+    for idx in xrange(len(node.outputs)):
+        if idx not in out_map:
+            for jdx in xrange(idx + 1, len(node.outputs)):
+                if (ts[idx] == ts[jdx] and
+                    fs[idx] == fs[jdx] and
+                    jdx not in out_map):
+                    out_map[jdx] = idx
+
+    if len(out_map.keys()) == 0:
+        return False
+
+    nw_ts = []
+    nw_fs = []
+    inv_map = {}
+    pos = 0
+    for idx in xrange(len(node.outputs)):
+        if idx not in out_map:
+            inv_map[idx] = pos
+            pos = pos + 1
+            nw_ts.append(ts[idx])
+            nw_fs.append(fs[idx])
+
+    new_ifelse = IfElse(n_outs=len(nw_ts),
+                        as_view=op.as_view,
+                        gpu=op.gpu,
+                        name=op.name)
+
+    new_ins = [node.inputs[0]] + nw_ts + nw_fs
+    new_outs = new_ifelse.make_node(*new_ins).outputs
+
+    rval = []
+    for idx in xrange(len(node.outputs)):
+        if idx in out_map.keys():
+            rval += [new_outs[inv_map[out_map[idx]]]]
+        else:
+            rval += [new_outs[inv_map[idx]]]
+
+    return rval
+
+
+@gof.local_optimizer([None])
+def cond_merge_random_op(main_node):
+    if isinstance(main_node.op, IfElse):
+        return False
+
+    all_inp_nodes = set()
+    for inp in main_node.inputs:
+        all_inp_nodes.add(inp.owner)
+    cond_nodes = [x for x in list(all_inp_nodes)
+                        if x and isinstance(x.op, IfElse)]
+
+    if len(cond_nodes) < 2:
+        return False
+
+    merging_node = cond_nodes[0]
+    for proposal in cond_nodes[1:]:
+        if (proposal.inputs[0] == merging_node.inputs[0] and
+            not find_up(proposal, merging_node) and
+            not find_up(merging_node, proposal)):
+            # Create a list of replacements for proposal
+            mn_ts = merging_node.inputs[1:][:merging_node.op.n_outs]
+            mn_fs = merging_node.inputs[1:][merging_node.op.n_outs:]
+            pl_ts = proposal.inputs[1:][:proposal.op.n_outs]
+            pl_fs = proposal.inputs[1:][proposal.op.n_outs:]
+            new_ins = ([merging_node.inputs[0]] +
+                       mn_ts + pl_ts + mn_fs + pl_fs)
+            mn_name = '?'
+            if merging_node.op.name:
+                mn_name = merging_node.op.name
+            pl_name = '?'
+            mn_n_ts = len(mn_ts)
+            mn_n_fs = len(mn_fs)
+            if proposal.op.name:
+                pl_name = proposal.op.name
+            new_ifelse = IfElse(
+                n_outs=len(mn_ts + pl_ts),
+                as_view=False,
+                gpu=False,
+                name=mn_name + '&' + pl_name)
+            new_outs = new_ifelse.make_node(*new_ins).outputs
+            old_outs = []
+            if type(merging_node.outputs) not in (list, tuple):
+                old_outs += [merging_node.outputs]
+            else:
+                old_outs += merging_node.outputs
+            if type(proposal.outputs) not in (list, tuple):
+                old_outs += [proposal.outputs]
+            else:
+                old_outs += proposal.outputs
+            pairs = zip(old_outs, new_outs)
+            main_outs = clone(main_node.outputs, replace=pairs)
+            return main_outs
+
+
+# XXX: Optimizations commented pending further debugging (certain optimizations
+# make computation less lazy than it should be currently).
+#
+# pushout_equilibrium = gof.EquilibriumDB()
+#
+# XXX: This optimization doesn't seem to exist anymore?
+# pushout_equilibrium.register("cond_lift_single_if",
+#                              opt.in2out(cond_lift_single_if,
+#                                         ignore_newtrees=True),
+#                              'fast_run', 'ifelse')
+#
+# pushout_equilibrium.register("cond_merge_random_op",
+#                              opt.in2out(cond_merge_random_op,
+#                                         ignore_newtrees=True),
+#                              'fast_run', 'ifelse')
+#
+#
+# pushout_equilibrium.register("ifelse_merge",
+#                              gof.MergeOptimizer(skip_const_merge=False),
+#                              'fast_run', 'ifelse')
+#
+# pushout_equilibrium.register("ifelse_remove_identical_inside",
+#                              opt.in2out(cond_remove_identical,
+#                                         ignore_newtrees=True),
+#                              'fast_run', 'ifelse')
+#
+# pushout_equilibrium.register('ifelse_sameCondTrue_inside',
+#                              opt.in2out(cond_merge_ifs_true,
+#                                         ignore_newtrees=True),
+#                              'fast_run', 'ifelse')
+#
+# pushout_equilibrium.register('ifelse_sameCondFalse_inside',
+#                              opt.in2out(cond_merge_ifs_false,
+#                                         ignore_newtrees=True),
+#                              'fast_run', 'ifelse')
+#
+# ifelse_seqopt.register('ifelse_condPushOut_equilibrium',
+#                        pushout_equilibrium,
+#                        1, 'fast_run', 'ifelse')
+#
+# ifelse_seqopt.register('merge_nodes_1',
+#                        gof.MergeOptimizer(skip_const_merge=False),
+#                        2, 'fast_run', 'ifelse')
+#
+#
+# ifelse_seqopt.register('ifelse_sameCondTrue',
+#                        opt.in2out(cond_merge_ifs_true,
+#                                   ignore_newtrees=True),
+#                        3, 'fast_run', 'ifelse')
+#
+#
+# ifelse_seqopt.register('ifelse_sameCondFalse',
+#                        opt.in2out(cond_merge_ifs_false,
+#                                   ignore_newtrees=True),
+#                        4, 'fast_run', 'ifelse')
+#
+#
+# ifelse_seqopt.register('ifelse_removeIdenetical',
+#                        opt.in2out(cond_remove_identical,
+#                                   ignore_newtrees=True),
+#                        7, 'fast_run', 'ifelse')

theano/lazycond.py

-"""
-IfElse is an Op that works with the LazyLinker to support conditional graph evaluation.
-
-:TODO: Add text to library documentation describing the IfElse Op.
-"""
-from copy import deepcopy
-import logging
-
-from theano.gof import PureOp, Apply, generic, Container
-
-import theano.tensor
-import gof
-
-from compile import optdb
-from tensor import opt
-
-_logger = logging.getLogger('theano.lazycond')
-
-@gof.local_optimizer([None])
-def ifelse_make_inplace(node):
-    op = node.op
-    if isinstance(op, IfElse) and not op.as_view :
-        _logger.debug('ifelse_make_inplace applied')
-        return IfElse(as_view = True,
-                    gpu = op.gpu, name=op.name).make_node(*node.inputs).outputs
-    return False
-
-optdb.register('ifelse_make_inplace', opt.in2out(ifelse_make_inplace,
-    ignore_newtrees=True), 95, 'fast_run', 'inplace')
-
-
-class IfElse(PureOp):
-    """
-    Op that works with LazyLinker to support conditional graph evaluation.
-    
-    Example usage:
-
-        ``rval = ifelse(tf, rval_if_true, rval_if_false)``
-
-    :type tf: symbolic tensor
-    :param tf: boolean variable representing a condition
-    :type rval_if_true: symbolic tensor
-    :param rval_if_false: symbolic variable to compute if tf is True
-    :type rval_if_false: symbolic tensor
-    :param rval_if_false: symbolic variable to compute if tf is False
-    :return: tensor corresponding to rval_if_true if tf is True or
-             rval_if_false if tf is False
-
-    While the switch function computes both values (rval_if_true and rval_if_false),
-    the ifelse op only computes one e.g rval_if_true is computed if tf is True.
-    
-    :note:
-        Other Linkers (ALL other linkers right now) are INCOMPATIBLE with this
-        Op, they will produce functions that FAIL TO EXECUTE.
-
-    """
-
-    def __init__(self, as_view=False, gpu = False, name = None):
-        if as_view:
-            # check destroyhandler and others to ensure that a view_map with
-            # multiple inputs can work
-            view_map = {}
-            view_map[0] = [1]
-            self.view_map = view_map
-            #raise NotImplementedError('IfElse must copy for now')
-        else:
-            self.view_map = {}
-        self.as_view=as_view
-        self.gpu = gpu
-        self.name = name
-
-    def __eq__(self, other):
-        return (type(self)==type(other) and
-                self.as_view == other.as_view and
-                # view_map included in as_view
-                #self.view_map == other.view_map and
-                self.gpu == other.gpu and
-                self.name == other.name)
-
-    def __hash__(self):
-        return (hash(type(self)) ^
-                # view_map included in as_view
-                # and dict are not hashable
-                #hash(self.view_map) ^
-                hash(self.as_view) ^
-                hash(self.gpu) ^
-                hash(self.name))
-
-    def make_node(self, c, t, f):
-        if t.type != f.type:
-            raise TypeError(
-                    'IfElse requires same types for true and false args',
-                    (t.type, f.type))
-        return Apply(self, [c,t,f], [t.type()])
-
-
-    def make_thunk(self, node, storage_map, compute_map, no_recycling):
-        outtype = node.outputs[0].type
-        c,t,f = node.inputs
-        output = node.outputs[0]
-        def thunk():
-            if not compute_map[c][0]:
-                return [0]
-            else:
-                truthval = storage_map[c][0]
-                if truthval:
-                    if not compute_map[t][0]:
-                        return [1]
-                    else:
-                        compute_map[output][0]=1
-                        if self.as_view:
-                            oval = outtype.filter(storage_map[t][0])
-                        else:
-                            oval = outtype.filter(
-                                    deepcopy(storage_map[t][0]))
-                        storage_map[output][0] = oval
-                        return []
-                else:
-                    if not compute_map[f][0]:
-                        return [2]
-                    else:
-                        # can't view both outputs unless destroyhandler
-                        # improves
-                        compute_map[output][0]=1
-                        oval = outtype.filter(
-                                deepcopy(storage_map[f][0]))
-                        storage_map[output][0]=oval
-                        return []
-        thunk.lazy = True
-        thunk.inputs  = [storage_map[v] for v in node.inputs]
-        thunk.outputs = [storage_map[v] for v in node.outputs]
-
-        return thunk
-
-ifelse = IfElse()

theano/tests/test_ifelse.py

+"""
+ Tests fof the lazy conditiona
+"""
+
+__docformat__ = 'restructedtext en'
+__authors__ = ("Razvan Pascanu ")
+__copyright__ = "(c) 2010, Universite de Montreal"
+__contact__ = "Razvan Pascanu <r.pascanu@gmail>"
+
+import unittest
+import numpy
+from nose.plugins.skip import SkipTest
+
+import theano
+from theano import tensor
+from theano.ifelse import IfElse, ifelse
+from theano.tests  import unittest_tools as utt
+
+
+class test_ifelse(unittest.TestCase):
+    def test_lazy_if(self):
+        # Tests that lazy if works .. even if the two results have different
+        # shapes but the same type (i.e. both vectors, or matrices or
+        # whatnot of same dtype
+        x = tensor.vector('x')
+        y = tensor.vector('y')
+        c = tensor.iscalar('c')
+        f = theano.function([c, x, y], ifelse(c, x, y))
+        rng = numpy.random.RandomState(utt.fetch_seed())
+
+        xlen = rng.randint(200)
+        ylen = rng.randint(200)
+
+        vx = numpy.asarray(rng.uniform(size=(xlen,)), theano.config.floatX)
+        vy = numpy.asarray(rng.uniform(size=(ylen,)), theano.config.floatX)
+
+        assert numpy.allclose(vx, f(1, vx, vy))
+        assert numpy.allclose(vy, f(0, vx, vy))
+
+    def test_lazy_if_on_generics(self):
+        x = theano.generic()
+        y = theano.generic()
+        c = tensor.iscalar('c')
+        f = theano.function([c, x, y], ifelse(c, x, y))
+
+        vx = ['testX']
+        vy = ['testY']
+        assert f(1, vx, vy) == vx
+        assert f(0, vx, vy) == vy
+
+    def test_grad_lazy_if(self):
+        # Tests that we can compute the gradients through lazy if
+        x = tensor.vector('x')
+        y = tensor.vector('y')
+        c = tensor.iscalar('c')
+        z = ifelse(c, x, y)
+        gx, gy = tensor.grad(z.sum(), [x, y])
+
+        f = theano.function([c, x, y], [gx, gy])
+        rng = numpy.random.RandomState(utt.fetch_seed())
+
+        xlen = rng.randint(200)
+        ylen = rng.randint(200)
+
+        vx = numpy.asarray(rng.uniform(size=(xlen,)), theano.config.floatX)
+        vy = numpy.asarray(rng.uniform(size=(ylen,)), theano.config.floatX)
+        gx0, gy0 = f(1, vx, vy)
+        assert numpy.allclose(gx0.shape, vx.shape)
+        assert numpy.allclose(gy0.shape, vy.shape)
+        assert numpy.all(gx0 == 1.)
+        assert numpy.all(gy0 == 0.)
+
+        gx0, gy0 = f(0, vx, vy)
+        assert numpy.allclose(gx0.shape, vx.shape)
+        assert numpy.allclose(gy0.shape, vy.shape)
+        assert numpy.all(gx0 == 0.)
+        assert numpy.all(gy0 == 1.)
+
+    def test_merge(self):
+        raise SkipTest("Optimization temporarily disabled")
+        x = tensor.vector('x')
+        y = tensor.vector('y')
+        c = tensor.iscalar('c')
+        z1 = ifelse(c, x + 1, y + 1)
+        z2 = ifelse(c, x + 2, y + 2)
+        z = z1 + z2
+        f = theano.function([c, x, y], z)
+        assert len([x for x in f.maker.env.toposort()
+                    if isinstance(x.op, IfElse)]) == 1
+
+    def test_remove_useless_inputs1(self):
+        raise SkipTest("Optimization temporarily disabled")
+        x = tensor.vector('x')
+        y = tensor.vector('y')
+        c = tensor.iscalar('c')
+        z = ifelse(c, (x, x), (y, y))
+        f = theano.function([c, x, y], z)
+
+        ifnode = [x for x in f.maker.env.toposort()
+                  if isinstance(x.op, IfElse)][0]
+        assert len(ifnode.inputs) == 3
+
+    def test_remove_useless_inputs2(self):
+        raise SkipTest("Optimization temporarily disabled")
+        x1 = tensor.vector('x1')
+        x2 = tensor.vector('x2')
+        y1 = tensor.vector('y1')
+        y2 = tensor.vector('y2')
+        c = tensor.iscalar('c')
+        z = ifelse(c, (x1, x1, x1, x2, x2), (y1, y1, y2, y2, y2))
+        f = theano.function([c, x1, x2, y1, y2], z)
+
+        ifnode = [x for x in f.maker.env.toposort()
+                  if isinstance(x.op, IfElse)][0]
+        assert len(ifnode.outputs) == 3
+
+    def test_pushout1(self):
+        raise SkipTest("Optimization temporarily disabled")
+        x1 = tensor.scalar('x1')
+        x2 = tensor.scalar('x2')
+        y1 = tensor.scalar('y1')
+        y2 = tensor.scalar('y2')
+        w1 = tensor.scalar('w1')
+        w2 = tensor.scalar('w2')
+        c = tensor.iscalar('c')
+        x, y = ifelse(c, (x1, y1), (x2, y2), name='f1')
+        z = ifelse(c, w1, w2, name='f2')
+        out = x * z * y
+
+        f = theano.function([x1, x2, y1, y2, w1, w2, c], out,
+                            allow_input_downcast=True)
+        assert isinstance(f.maker.env.toposort()[-1].op, IfElse)
+        rng = numpy.random.RandomState(utt.fetch_seed())
+        vx1 = rng.uniform()
+        vx2 = rng.uniform()
+        vy1 = rng.uniform()
+        vy2 = rng.uniform()
+        vw1 = rng.uniform()
+        vw2 = rng.uniform()
+
+        assert numpy.allclose(f(vx1, vx2, vy1, vy2, vw1, vw2, 1),
+                              vx1 * vy1 * vw1)
+        assert numpy.allclose(f(vx1, vx2, vy1, vy2, vw1, vw2, 0),
+                              vx2 * vy2 * vw2)
+
+    def test_pushout3(self):
+        raise SkipTest("Optimization temporarily disabled")
+        x1 = tensor.scalar('x1')
+        y1 = tensor.scalar('x2')
+        y2 = tensor.scalar('y2')
+        c = tensor.iscalar('c')
+        two = numpy.asarray(2, dtype=theano.config.floatX)
+        x, y = ifelse(c, (x1, y1), (two, y2), name='f1')
+        o3 = numpy.asarray(0.3, dtype=theano.config.floatX)
+        o2 = numpy.asarray(0.2, dtype=theano.config.floatX)
+        z = ifelse(c, o3, o2, name='f2')
+        out = x * z * y
+
+        f = theano.function([x1, y1, y2, c], out,
+                            allow_input_downcast=True)
+        assert isinstance(f.maker.env.toposort()[-1].op, IfElse)
+        rng = numpy.random.RandomState(utt.fetch_seed())
+        vx1 = rng.uniform()
+        vy1 = rng.uniform()
+        vy2 = rng.uniform()
+
+        assert numpy.allclose(f(vx1, vy1, vy2, 1), vx1 * vy1 * 0.3)
+        assert numpy.allclose(f(vx1, vy1, vy2, 0), 2 * vy2 * 0.2)
+
+    def test_pushout2(self):
+        raise SkipTest("Optimization temporarily disabled")
+        x1 = tensor.scalar('x1')
+        x2 = tensor.scalar('x2')
+        y1 = tensor.scalar('y1')
+        y2 = tensor.scalar('y2')
+        w1 = tensor.scalar('w1')
+        w2 = tensor.scalar('w2')
+        c = tensor.iscalar('c')
+        x, y = ifelse(c, (x1, y1), (x2, y2), name='f1')
+        z = ifelse(x > y, w1, w2, name='f2')
+        out = x * z * y
+
+        f = theano.function([x1, x2, y1, y2, w1, w2, c], out,
+                            allow_input_downcast=True)
+        assert isinstance(f.maker.env.toposort()[-1].op, IfElse)
+        rng = numpy.random.RandomState(utt.fetch_seed())
+        vx1 = rng.uniform()
+        vx2 = rng.uniform()
+        vy1 = rng.uniform()
+        vy2 = rng.uniform()
+        vw1 = rng.uniform()
+        vw2 = rng.uniform()
+        if vx1 > vy1:
+            vw = vw1
+        else:
+            vw = vw2
+        assert numpy.allclose(f(vx1, vx2, vy1, vy2, vw1, vw2, 1),
+                              vx1 * vy1 * vw)
+
+        if vx2 > vy2:
+            vw = vw1
+        else:
+            vw = vw2
+        assert numpy.allclose(f(vx1, vx2, vy1, vy2, vw1, vw2, 0),
+                              vx2 * vy2 * vw)
+
+    def test_merge_ifs_true_false(self):
+        raise SkipTest("Optimization temporarily disabled")
+        x1 = tensor.scalar('x1')
+        x2 = tensor.scalar('x2')
+        y1 = tensor.scalar('y1')
+        y2 = tensor.scalar('y2')
+        w1 = tensor.scalar('w1')
+        w2 = tensor.scalar('w2')
+        c = tensor.iscalar('c')
+
+        out = ifelse(c,
+            ifelse(c, x1, x2) + ifelse(c, y1, y2) + w1,
+            ifelse(c, x1, x2) + ifelse(c, y1, y2) + w2)
+        f = theano.function([x1, x2, y1, y2, w1, w2, c], out,
+                            allow_input_downcast=True)
+        assert len([x for x in f.maker.env.toposort()
+                if isinstance(x.op, IfElse)]) == 1
+
+        rng = numpy.random.RandomState(utt.fetch_seed())
+        vx1 = rng.uniform()
+        vx2 = rng.uniform()
+        vy1 = rng.uniform()
+        vy2 = rng.uniform()
+        vw1 = rng.uniform()
+        vw2 = rng.uniform()
+        assert numpy.allclose(f(vx1, vx2, vy1, vy2, vw1, vw2, 1),
+                              vx1 + vy1 + vw1)
+        assert numpy.allclose(f(vx1, vx2, vy1, vy2, vw1, vw2, 0),
+                              vx2 + vy2 + vw2)
+
+
+if __name__ == '__main__':
+    print ' Use nosetests to run these tests '