@gof.local_optimizer(acceptable_ops)
def ifelse_lift_single_if_through_acceptable_ops(main_node):
    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(*true_ins, **dict(return_list=True))
    false_eval = mop(*false_ins, **dict(return_list=True))
    # 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, return_list=True)
    return nw_outs
This optimization lifts up certain ifelse instances.
dataset/ETHPy150Open rizar/attention-lvcsr/libs/Theano/theano/ifelse.py/ifelse_lift_single_if_through_acceptable_ops