from __future__ import absolute_import, print_function, division
"""
Tensor optimizations addressing the ops in basic.py.
"""
# TODO: intelligent merge for mul/add
# TODO: 0*x -> 0
from collections import defaultdict
import logging
import itertools
import operator
import sys
import time
import traceback
import warnings
import numpy
from six import integer_types, iteritems
from six.moves import reduce, xrange
import theano
from theano import gof
from theano.compat import izip
from theano.gof import opt, InconsistencyError, TopoOptimizer, graph
from theano.gof import Variable, Constant
from theano.gof.opt import copy_stack_trace, in2out
from theano.gof.utils import MethodNotDefined
from theano.gradient import DisconnectedType
from theano.configparser import config
from theano.tensor.elemwise import Elemwise, DimShuffle
from theano.tensor.subtensor import (get_idx_list, get_canonical_form_slice,
                                     Subtensor, IncSubtensor, make_constant,
                                     AdvancedIncSubtensor1,
                                     AdvancedIncSubtensor,
                                     AdvancedSubtensor1,
                                     advanced_subtensor,
                                     advanced_subtensor1,
                                     advanced_inc_subtensor1)
from theano import scalar
from theano.scalar import basic
from theano.tensor import basic as T
from theano import compile  # to register the optimizer built by this file
from theano.compile.ops import Shape, Shape_i
from theano.tensor.type import (values_eq_approx_remove_inf,
                                values_eq_approx_remove_nan,
                                values_eq_approx_remove_inf_nan)
from theano.gof.opt import (Optimizer, pre_constant_merge,
                            pre_greedy_local_optimizer)
from theano.gof import toolbox
from theano.tensor.basic import (Alloc, get_scalar_constant_value, ShapeError,
                                 extract_constant, NotScalarConstantError,
                                 Reshape)
from six import StringIO
_logger = logging.getLogger('theano.tensor.opt')
# Utilities
def _fill_chain(new_out, orig_inputs):
    for i in orig_inputs:
        new_out = T.fill(i, new_out)
    return [new_out]
def encompasses_broadcastable(b1, b2):
    """
    Parameters
    ----------
    b1
        The broadcastable attribute of a tensor type.
    b2
        The broadcastable attribute of a tensor type.
    Returns
    -------
    bool
        True if the broadcastable patterns b1 and b2 are such that b2 is
        broadcasted to b1's shape and not the opposite.
    """
    if len(b1) < len(b2):
        return False
    b1 = b1[-len(b2):]
    return not any(v1 and not v2 for v1, v2 in zip(b1, b2))
def merge_broadcastables(broadcastables):
    return [all(bcast) for bcast in zip(*broadcastables)]
def scalarconsts_rest(inputs, elemwise=True, only_process_constants=False):
    """Partition a list of variables into two kinds:
    scalar constants, and the rest."""
    consts = []
    origconsts = []
    nonconsts = []
    for i in inputs:
        try:
            v = get_scalar_constant_value(i, elemwise=elemwise,
                                          only_process_constants=only_process_constants)
            consts.append(v)
            origconsts.append(i)
        except NotScalarConstantError:
            nonconsts.append(i)
    return consts, origconsts, nonconsts
def broadcast_like(value, template, fgraph, dtype=None):
    """
    Return a Variable with the same shape and dtype as the template,
    filled by broadcasting value through it. `value` will be cast as
    necessary.
    """
    value = T.as_tensor_variable(value)
    if value.type == template.type:
        return value
    if template not in fgraph.variables:
        raise NotImplementedError('broadcast_like currently requires the '
                                  'template Variable to be in the fgraph already')
    if dtype is None:
        dtype = template.dtype
    value = T.cast(value, dtype)
    if value.type == template.type:
        return value
    if hasattr(fgraph, 'shape_feature'):
        new_shape = fgraph.shape_feature.shape_of[template]
    else:
        new_shape = template.shape
    rval = T.alloc(value, *new_shape)
    # the template may have 1s in its shape without being broadcastable
    if rval.broadcastable != template.broadcastable:
        rval = T.unbroadcast(rval, *[i for i in xrange(rval.ndim)
                                     if rval.broadcastable[i] and
                                     not template.broadcastable[i]])
    assert rval.type.dtype == dtype
    if rval.type.broadcastable != template.broadcastable:
        raise AssertionError("rval.type.broadcastable is " +
                             str(rval.type.broadcastable) +
                             " but template.broadcastable is" +
                             str(template.broadcastable))
    return rval
class InplaceElemwiseOptimizer(Optimizer):
    """
    We parametrise it to make it work for Elemwise and GpuElemwise op.
    """
    def __init__(self, OP):
        self.op = OP
    def add_requirements(self, fgraph):
        fgraph.attach_feature(theano.gof.destroyhandler.DestroyHandler())
    @staticmethod
    def print_profile(stream, prof, level=0):
        blanc = ('    ' * level)
        print(blanc, "InplaceElemwiseOptimizer ", prof['opt'].op, file=stream)
        for k in ['node_before',
                  'nb_call_replace',
                  'nb_call_validate',
                  'nb_inconsistent']:
            print(blanc, k, prof[k], file=stream)
        ndim = prof['ndim']
        if ndim:
            print(blanc, "ndim", "nb", file=stream)
            for n in sorted(ndim.keys()):
                print(blanc, n, ndim[n], file=stream)
    def apply(self, fgraph):
        """
        Usage: InplaceElemwiseOptimizer(op).optimize(fgraph)
        Attempts to replace all Broadcast ops by versions of them
        that operate inplace. It operates greedily: for each Broadcast
        Op that is encountered, for each output, tries each input to
        see if it can operate inplace on that input. If so, makes the
        change and go to the next output or Broadcast Op.
        Examples
        --------
            `x + y + z -> x += y += z`
            `(x + y) * (x * y) -> (x += y) *= (x * y) or (x + y) *= (x *= y)`
        """
        # We should not validate too often as this takes too much time to
        # execute!
        # It is the _dfs_toposort() fct in theano/gof/destroyhandler.py
        # that takes so much time.
        # Should we try to use another lib that does toposort?
        #   igraph: http://igraph.sourceforge.net/
        #   networkx: https://networkx.lanl.gov/
        # Should we try to use cython?
        #   Compiling only that fct is not enough, should we try to add the
        #   deque class too?
        #   And init the deque and other list to an upper bound number of
        #   elements?
        # Maybe Theano should do online toposort as in
        #   http://code.google.com/p/acyclic
        #
        # The next longest optimizer is the canonizer phase.
        # Then I think it is the [io_?]toposort (need to validate) so check if
        # the solution is also applicable there.
        # We execute `validate` after this number of change.
        prof = {'opt': self,
                'node_before': len(fgraph.apply_nodes),
                'nb_call_replace': 0,
                'nb_call_validate': 0,
                'nb_inconsistent': 0,
                'ndim': defaultdict(lambda: 0)}
        check_each_change = config.tensor.insert_inplace_optimizer_validate_nb
        if check_each_change == -1:
            if len(fgraph.apply_nodes) > 500:
                check_each_change = 10
            else:
                check_each_change = 1
        nb_change_no_validate = 0
        chk = fgraph.checkpoint()
        if fgraph.update_mapping:
            update_outs = [fgraph.outputs[i] for i in fgraph.update_mapping]
        else:
            update_outs = []
        protected_inputs = [
            f.protected for f in fgraph._features if
            isinstance(f, theano.compile.function_module.Supervisor)]
        protected_inputs = sum(protected_inputs, [])  # flatten the list
        protected_inputs.extend(fgraph.outputs)
        for node in list(graph.io_toposort(fgraph.inputs, fgraph.outputs)):
            op = node.op
            # gpuarray GpuElemwise inherit from Elemwise
            if not type(op) == self.op:
                continue
            # If big graph and the outputs are scalar, do not make it
            # inplace.
            if (check_each_change != 1 and
                # If multiple outputs, they must all have the same size,
                # so only check the first.
                    getattr(node.outputs[0].type, 'ndim', -1) == 0):
                continue
            if op.inplace_pattern:
                # Maybe this isn't needed anymore, but I don't want to
                # rish regression now. This case only happen if the
                # original node add already some inplace patter and we
                # still try to add more pattern.
                baseline = op.inplace_pattern
                candidate_outputs = [i for i in xrange(len(node.outputs))
                                     if i not in baseline]
                # node inputs that are Constant, already destroyed,
                # or fgraph protected inputs and fgraph outputs can't be used as
                # inplace target.
                # Remove here as faster.
                candidate_inputs = [i for i in xrange(len(node.inputs))
                                    if i not in baseline.values() and
                                    not isinstance(node.inputs[i], Constant) and
                                    # Is next line costly?
                                    not fgraph.destroyers(node.inputs[i]) and
                                    node.inputs[i] not in protected_inputs]
            else:
                baseline = []
                candidate_outputs = list(range(len(node.outputs)))
                # node inputs that are Constant, already destroyed,
                # fgraph protected inputs and fgraph outputs can't be used as inplace
                # target.
                # Remove here as faster.
                candidate_inputs = [i for i in xrange(len(node.inputs))
                                    if not isinstance(node.inputs[i], Constant) and
                                    not fgraph.destroyers(node.inputs[i]) and
                                    node.inputs[i] not in protected_inputs]
            verbose = False
            raised_warning = not verbose
            for candidate_output in candidate_outputs:
                # If the output of the node can be established as an update
                # output of the fgraph, visit the candidate_inputs in an order
                # that will improve the chances of making the node operate
                # inplace on the input it's meant to update
                candidate_out_var = node.outputs[candidate_output]
                sorted_candidate_inputs = candidate_inputs
                if candidate_out_var in update_outs:
                    # The candidate output is an update. Sort the
                    # variables in candidate_inputs in the following order:
                    # - Vars corresponding to the actual updated input
                    #   (best case scenario is for the node that procudes
                    #   an update to operate inplace on the variable to
                    #   update)
                    # - Vars computed inplace on the updates input (second
                    #   best scenario if for the node to work inplace on
                    #   a variable obtained by a chain of inplace on the
                    #   variable to update. In some cases, this will be
                    #   equivalent to operating inplace on the variable to
                    #   update)
                    # - Remaining variables
                    updated_inputs = []
                    for i, f_out in enumerate(fgraph.outputs):
                        if (f_out is candidate_out_var and i in fgraph.update_mapping):
                            updated_inp_idx = fgraph.update_mapping[i]
                            updated_inputs.append(fgraph.inputs[updated_inp_idx])
                    updated_vars = []
                    vars_from_inplace = []
                    other_vars = []
                    for inp_idx in candidate_inputs:
                        inp = node.inputs[inp_idx]
                        if inp in updated_inputs:
                            # the candidate input is the actual updated input
                            updated_vars.append(inp_idx)
                        elif (hasattr(fgraph, 'destroy_handler') and
                              inp.owner and
                              any([fgraph.destroy_handler.root_destroyer.get(up_inp, None) is inp.owner
                                   for up_inp in updated_inputs])):
                            # the candidate input is a variable computed
                            # inplace on the updated input via a sequence of
                            # one or more inplace operations
                            vars_from_inplace.append(inp_idx)
                        else:
                            other_vars.append(inp_idx)
                    sorted_candidate_inputs = (updated_vars +
                                               vars_from_inplace + other_vars)
                for candidate_input in sorted_candidate_inputs:
                    # remove inputs that don't have the same dtype as the output
                    if node.inputs[candidate_input].type != node.outputs[
                            candidate_output].type:
                        continue
                    inplace_pattern = dict(baseline)
                    inplace_pattern[candidate_output] = candidate_input
                    try:
                        if hasattr(op.scalar_op, "make_new_inplace"):
                            new_scal = op.scalar_op.make_new_inplace(
                                scalar.transfer_type(
                                    *[inplace_pattern.get(i, o.dtype)
                                      for i, o in enumerate(node.outputs)]))
                        else:
                            new_scal = op.scalar_op.__class__(
                                scalar.transfer_type(
                                    *[inplace_pattern.get(i, None)
                                      for i in xrange(len(node.outputs))]))
                        new_outputs = self.op(new_scal, inplace_pattern)(
                            *node.inputs, **dict(return_list=True))
                        new_node = new_outputs[0].owner
                        for r, new_r in zip(node.outputs, new_outputs):
                            prof['nb_call_replace'] += 1
                            fgraph.replace(r, new_r,
                                           reason="inplace_elemwise_optimizer")
                        nb_change_no_validate += 1
                        prof['ndim'][candidate_out_var.ndim] += 1
                        if nb_change_no_validate >= check_each_change:
                            prof['nb_call_validate'] += 1
                            fgraph.validate()
                            chk = fgraph.checkpoint()
                            nb_change_no_validate = 0
                    except (ValueError, InconsistencyError) as e:
                        prof['nb_inconsistent'] += 1
                        if check_each_change != 1 and not raised_warning:
                            print(("Some inplace optimization was not "
                                   "performed due to unexpected error:"),
                                  file=sys.stderr)
                            print(e, file=sys.stderr)
                            raised_warning = True
                        fgraph.revert(chk)
                        continue
                    candidate_inputs.remove(candidate_input)
                    node = new_node
                    baseline = inplace_pattern
                    break
        if nb_change_no_validate > 0:
            try:
                fgraph.validate()
            except Exception:
                if not raised_warning:
                    print(("Some inplace optimization was not "
                           "performed due to unexpected error"),
                          file=sys.stderr)
                fgraph.revert(chk)
        return prof
    def print_summary(self, stream=sys.stdout, level=0, depth=-1):
        print("%s%s (%s)" % (
            (' ' * level), self.__class__.__name__, self.op), file=stream)
        return inplace_elemwise_optimizer
inplace_elemwise_optimizer = InplaceElemwiseOptimizer(T.Elemwise)
compile.optdb.register('inplace_elemwise_opt', inplace_elemwise_optimizer, 75,
                       'inplace_opt',  # for historic reason
                       'inplace_elemwise_optimizer',
                       'fast_run', 'inplace')
def register_useless(lopt, *tags, **kwargs):
    if type(lopt) == str:
        def register(inner_lopt):
            return register_useless(inner_lopt, lopt, *tags, **kwargs)
        return register
    else:
        name = kwargs.pop('name', None) or lopt.__name__
        compile.mode.local_useless.register(name, lopt, 'last', 'fast_run',
                                            *tags, **kwargs)
        return lopt
def register_canonicalize(lopt, *tags, **kwargs):
    if type(lopt) == str:
        def register(inner_lopt):
            return register_canonicalize(inner_lopt, lopt, *tags, **kwargs)
        return register
    else:
        name = kwargs.pop('name', None) or lopt.__name__
        compile.optdb['canonicalize'].register(name, lopt, 'fast_run',
                                               *tags, **kwargs)
        return lopt
def register_stabilize(lopt, *tags, **kwargs):
    if type(lopt) == str:
        def register(inner_lopt):
            return register_stabilize(inner_lopt, lopt, *tags, **kwargs)
        return register
    else:
        name = kwargs.pop('name', None) or lopt.__name__
        compile.optdb['stabilize'].register(name, lopt, 'fast_run',
                                            *tags, **kwargs)
        return lopt
def register_specialize(lopt, *tags, **kwargs):
    if type(lopt) == str:
        def register(inner_lopt):
            return register_specialize(inner_lopt, lopt, *tags, **kwargs)
        return register
    else:
        name = kwargs.pop('name', None) or lopt.__name__
        compile.optdb['specialize'].register(name, lopt, 'fast_run',
                                             *tags, **kwargs)
        return lopt
def register_uncanonicalize(lopt, *tags, **kwargs):
    if type(lopt) == str:
        def register(inner_lopt):
            return register_uncanonicalize(inner_lopt, lopt, *tags, **kwargs)
        return register
    else:
        name = (kwargs and kwargs.pop('name', None)) or lopt.__name__
        compile.optdb['uncanonicalize'].register(name, lopt, 'fast_run', *tags,
                                                 **kwargs)
        return lopt
def register_specialize_device(lopt, *tags, **kwargs):
    if type(lopt) == str:
        def register(inner_lopt):
            return register_specialize_device(inner_lopt, lopt, *tags, **kwargs)
        return register
    else:
        name = (kwargs and kwargs.pop('name', None)) or lopt.__name__
        compile.optdb['specialize_device'].register(name, lopt, 'fast_run', *tags,
                                                    **kwargs)
        return lopt
#####################
# Dot optimizations #
#####################
@register_canonicalize
@register_stabilize
@gof.local_optimizer([T.Dot])
def local_0_dot_x(node):
    if not isinstance(node.op, T.Dot):
        return False
    x = node.inputs[0]
    y = node.inputs[1]
    replace = False
    try:
        if get_scalar_constant_value(x, only_process_constants=True) == 0:
            replace = True
    except NotScalarConstantError:
        pass
    try:
        if get_scalar_constant_value(y, only_process_constants=True) == 0:
            replace = True
    except NotScalarConstantError:
        pass
    if replace:
        constant_zero = T.constant(0, dtype=node.outputs[0].type.dtype)
        if x.ndim == 2 and y.ndim == 2:
            constant_zero = assert_(constant_zero,
                                    T.eq(x.shape[1], y.shape[0]))
            return [T.alloc(constant_zero, x.shape[0], y.shape[1])]
        elif x.ndim == 1 and y.ndim == 2:
            constant_zero = assert_(constant_zero,
                                    T.eq(x.shape[0], y.shape[0]))
            return [T.alloc(constant_zero, y.shape[1])]
        elif x.ndim == 2 and y.ndim == 1:
            constant_zero = assert_(constant_zero,
                                    T.eq(x.shape[1], y.shape[0]))
            return [T.alloc(constant_zero, x.shape[0])]
        elif x.ndim == 1 and y.ndim == 1:
            constant_zero = assert_(constant_zero,
                                    T.eq(x.shape[0], y.shape[0]))
            return [constant_zero]
        else:
            _logger.warning("Optimization Warning: "
                            "Optimization theano/opt.py:local_0_dot_x Found "
                            "that it could apply, but was not implemented "
                            "for dot product with these input types:\n"
                            "(%s, %s)",
                            x.type, y.type)
######################
# DimShuffle lifters #
######################
def apply_local_dimshuffle_lift(var):
    # return var
    # lift recursively
    if not var.owner:
        return var
    new = local_dimshuffle_lift.transform(var.owner)
    if new:
        return new[0]
    return var
# Checks for two types of useless dimshuffles:
#   1 - dimshuffle all dimensions in order.
#   2 - dimshuffle a broadcastable dimension.
def is_dimshuffle_useless(new_order, input):
    is_useless = True
    if len(new_order) == input.type.ndim:
        all_broadcastable_dims = [i for (i, is_broadcastable)
                                  in enumerate(input.type.broadcastable)
                                  if is_broadcastable] + ['x']
        for i in range(input.type.ndim):
            if (new_order[i] == i or
                    (i in all_broadcastable_dims and
                     new_order[i] in all_broadcastable_dims)):
                is_useless = True
            else:
                is_useless = False
                break
    else:
        is_useless = False
    return is_useless
@gof.local_optimizer([DimShuffle])
def local_dimshuffle_lift(node):
    """
    "Lifts" DimShuffle through Elemwise operations and merges
    consecutive DimShuffles. Basically, applies the following
    transformations on the whole graph:
    DimShuffle(Elemwise(x, y)) => Elemwise(DimShuffle(x), DimShuffle(y))
    DimShuffle(DimShuffle(x)) => DimShuffle(x)
    DimShuffle{0,1,...}(x) => x (when the dimshuffle do nothing)
    After this transform, clusters of Elemwise operations are
    void of DimShuffle operations.
    """
    op = node.op
    if not isinstance(op, DimShuffle):
        return False
    input = node.inputs[0]
    inode = input.owner
    new_order = op.new_order
    if inode and isinstance(inode.op, Elemwise) and (len(input.clients) == 1):
        # Don't use make_node to have tag.test_value set.
        new_inputs = []
        for inp in inode.inputs:
            new_inp = op.__class__(inp.type.broadcastable,
                                   op.new_order)(inp)
            new_inputs.append(apply_local_dimshuffle_lift(new_inp))
        copy_stack_trace(node.outputs[0], new_inputs)
        ret = inode.op(*new_inputs, **dict(return_list=True))
        return ret
    if inode and isinstance(inode.op, DimShuffle):
        new_order = [x == 'x' and 'x' or inode.op.new_order[x] for x in
                     new_order]
        input = inode.inputs[0]
    if is_dimshuffle_useless(new_order, input):
        return [input]
    elif inode and isinstance(inode.op, DimShuffle):
        ret = op.__class__(input.type.broadcastable, new_order)(input)
        ret = apply_local_dimshuffle_lift(ret)
        copy_stack_trace(node.outputs[0], ret)
        return [ret]
@register_canonicalize
@gof.local_optimizer([Reshape])
def local_useless_dimshuffle_in_reshape(node):
    """
    Removes useless DimShuffle operation inside Reshape:
      reshape(vector.dimshuffle('x', 0), shp) => reshape(vector, shp)
      reshape(matrix.dimshuffle('x', 0, 'x', 1), shp) => reshape(matrix, shp)
      reshape(row.dimshuffle(1, 'x'), shp) => reshape(row, shp)
      reshape(col.dimshuffle(0), shp) => reshape(col, shp)
    """
    op = node.op
    if not isinstance(op, Reshape):
        return False
    if not (node.inputs[0].owner is not None and
            isinstance(node.inputs[0].owner.op, DimShuffle)):
        return False
    new_order = node.inputs[0].owner.op.new_order
    input = node.inputs[0].owner.inputs[0]
    broadcastables = node.inputs[0].broadcastable
    new_order_of_nonbroadcast = []
    for i, bd in zip(new_order, broadcastables):
        if not bd:
            new_order_of_nonbroadcast.append(i)
    no_change_in_order = all(
        new_order_of_nonbroadcast[i] <= new_order_of_nonbroadcast[i + 1]
        for i in xrange(len(new_order_of_nonbroadcast) - 1))
    if no_change_in_order:
        shape = node.inputs[1]
        ret = op.__class__(node.outputs[0].ndim)(input, shape)
        copy_stack_trace(node.outputs[0], ret)
        return [ret]
@register_canonicalize
@gof.local_optimizer([DimShuffle])
def local_lift_transpose_through_dot(node):
    """
    dot(x,y).T -> dot(y.T, x.T)
    These optimizations "lift" (propagate towards the inputs) DimShuffle
    through dot product.  It allows to put the graph in a more standard shape,
    and to later merge consecutive DimShuffles.
    The transformation should be apply whether or not the transpose is
    inplace.  The newly-introduced transpositions are not inplace, this will
    be taken care of in a later optimization phase.
    """
    if not (isinstance(node.op, T.DimShuffle) and node.op.new_order == (1, 0)):
        return False
    if not (node.inputs[0].owner and
            isinstance(node.inputs[0].owner.op, T.Dot)):
        return False
    x, y = node.inputs[0].owner.inputs
    if x.ndim == y.ndim == 2:
        # Output is dot product of transposed inputs in reverse order
        ret = [T.dot(y.T, x.T)]
        # Copy over stack trace to output from result of dot-product
        copy_stack_trace(node.inputs[0], ret)
        return ret
register_canonicalize(local_dimshuffle_lift)
register_specialize(local_dimshuffle_lift)
######################
# Casting operations #
######################
@register_canonicalize
@register_specialize
@gof.local_optimizer([T.TensorFromScalar])
def local_tensor_scalar_tensor(node):
    '''tensor_from_scalar(scalar_from_tensor(x)) -> x'''
    if isinstance(node.op, T.TensorFromScalar):
        s = node.inputs[0]
        if s.owner and isinstance(s.owner.op, T.ScalarFromTensor):
            t = s.owner.inputs[0]
            # We don't need to copy over any stack traces here
            return [t]
@register_canonicalize
@register_specialize
@gof.local_optimizer([T.ScalarFromTensor])
def local_scalar_tensor_scalar(node):
    '''scalar_from_tensor(tensor_from_scalar(x)) -> x'''
    if isinstance(node.op, T.ScalarFromTensor):
        t = node.inputs[0]
        if t.owner and isinstance(t.owner.op, T.TensorFromScalar):
            s = t.owner.inputs[0]
            # We don't need to copy over any stack traces here
            return [s]
#####################################
# ShapeFeature, Shape optimizations
#####################################
class MakeVector(T.Op):
    """Concatenate a number of scalars together into a vector.
    This is a simple version of stack() that introduces far less cruft
    into the graph. Should work with 0 inputs. The constant_folding
    optimization will remove it.
    """
    __props__ = ("dtype",)
    def __init__(self, dtype='int64'):
        self.dtype = dtype
    def make_node(self, *inputs):
        inputs = list(map(T.as_tensor_variable, inputs))
        if (not all(a.type == inputs[0].type for a in inputs) or
                (len(inputs) > 0 and inputs[0].dtype != self.dtype)):
            dtype = theano.scalar.upcast(self.dtype, *[i.dtype for i in inputs])
            # upcast the input to the determined dtype,
            # but don't downcast anything
            assert dtype == self.dtype, (
                "The upcast of the inputs to MakeVector should match the "
                "dtype given in __init__.")
            if not all(self.dtype == T.cast(i, dtype=dtype).dtype
                       for i in inputs):
                raise TypeError("MakeVector.make_node expected inputs"
                                " upcastable to %s. got %s" %
                                (self.dtype, str([i.dtype for i in inputs])))
            inputs = [T.cast(i, dtype=dtype) for i in inputs]
        assert all(self.dtype == a.dtype for a in inputs)
        assert all(a.ndim == 0 for a in inputs)
        if inputs:
            dtype = inputs[0].type.dtype
        else:
            dtype = self.dtype
        # bcastable = (len(inputs) == 1)
        bcastable = False
        otype = T.TensorType(broadcastable=(bcastable,), dtype=dtype)
        return T.Apply(self, inputs, [otype()])
    def perform(self, node, inputs, out_):
        out, = out_
        # not calling theano._asarray as optimization
        if (out[0] is None) or (out[0].size != len(inputs)):
            out[0] = theano._asarray(inputs, dtype=node.outputs[0].dtype)
        else:
            # assume that out has correct dtype. there is no cheap way to check
            out[0][...] = inputs
    def c_code_cache_version(self):
        return (2,)
    def c_code(self, node, name, inp, out_, sub):
        out, = out_
        # Shouldn't use PyArray_TYPE(inp[0]) for the dtype
        # when len(inp) == 0 (we need to support this case.
        # So there will be (1 * nb_dtype) + ((nb len(inp) - 1 ))
        # different c code with the following algo
        out_shape = len(inp)
        out_num = numpy.dtype(node.outputs[0].dtype).num
        # don't use dtype_%(out)s as when check_input=False, it isn't defined.
        out_dtype = node.outputs[0].type.dtype_specs()[1]
        if len(inp) > 0:
            assert self.dtype == node.inputs[0].dtype
            out_num = 'PyArray_TYPE(%s)' % inp[0]
        ret = """
        npy_intp dims[1];
        dims[0] = %(out_shape)s;
        if(!%(out)s || PyArray_DIMS(%(out)s)[0] != %(out_shape)s){
            Py_XDECREF(%(out)s);
            %(out)s = (PyArrayObject*)PyArray_EMPTY(1, dims, %(out_num)s, 0);
        }
        """ % locals()
        for idx, i in enumerate(inp):
            ret += """
            *((%(out_dtype)s *)PyArray_GETPTR1(%(out)s, %(idx)s)) = *((%(out_dtype)s *) PyArray_DATA(%(i)s));
            """ % locals()
        return ret
    def infer_shape(self, node, ishapes):
        return [(len(ishapes),)]
    def grad(self, inputs, output_gradients):
        # If the output is of an integer dtype, no gradient shall pass
        if self.dtype in theano.tensor.discrete_dtypes:
            return [ipt.zeros_like().astype(theano.config.floatX)
                    for ipt in inputs]
        grads = []
        for i, inp in enumerate(inputs):
            grads.append(output_gradients[0][i])
        return grads
    def R_op(self, inputs, eval_points):
        if None in eval_points:
            return [None]
        return self.make_node(*eval_points).outputs
make_vector = MakeVector()
class MakeVectorPrinter:
    def process(self, r, pstate):
        if r.owner is None:
            raise TypeError("Can only print make_vector.")
        elif isinstance(r.owner.op, MakeVector):
            old_precedence = getattr(pstate, 'precedence', None)
            try:
                pstate.precedence = 1000
                s = [pstate.pprinter.process(input)
                     for input in r.owner.inputs]
            finally:
                pstate.precedence = old_precedence
            return "[%s]" % ", ".join(s)
        else:
            raise TypeError("Can only print make_vector.")
T.pprint.assign(MakeVector, MakeVectorPrinter())
class ShapeFeature(object):
    """Graph optimizer for removing all calls to shape().
    This optimizer replaces all Shapes and Subtensors of Shapes with
    Shape_i and MakeVector Ops.
    This optimizer has several goals:
    1. to 'lift' Shapes to as close to the inputs as possible.
    2. to infer the shape of every node in the graph in terms of the
       input shapes.
    3. remove all fills (T.second, T.fill) from the graph
    Lifting shapes as close to the inputs as possible is important for
    canonicalization because it is very bad form to have to compute
    something just to know how big it will be.  Firstly, it is a waste
    of time to compute such outputs.  But it is important to get rid
    of these outputs as early as possible in the compilation process
    because the extra computations make it appear as if many internal
    graph nodes have multiple clients.  Many optimizations refuse to
    work on nodes with multiple clients.
    Lifting is done by using an `<Op>.infer_shape` function if one is
    present, or else using a conservative default.  An Op that
    supports shape-lifting should define a infer_shape(self, node,
    input_shapes) function.  The argument input_shapes is a tuple of
    tuples... there is an interior tuple for each input to the node.
    The tuple has as many elements as dimensions.  The element in
    position i of tuple j represents the i'th shape component of the
    j'th input.  The function should return a tuple of tuples.  One
    output tuple for each node.output.  Again, the i'th element of the
    j'th output tuple represents the output[j].shape[i] of the
    function.  If an output is not a TensorType, then None should be
    returned instead of a tuple for that output.
    For example the infer_shape for a matrix-matrix product would accept
    input_shapes=((x0,x1), (y0,y1)) and return ((x0, y1),).
    Inferring the shape of internal nodes in the graph is important
    for doing size-driven optimizations.  If we know how big various
    intermediate results will be, we can estimate the cost of many Ops
    accurately, and generate c-code that is specific [e.g. unrolled]
    to particular sizes.
    In cases where you cannot figure out the shape, raise a ShapeError.
    Notes
    -----
    Right now there is only the ConvOp that could really take
    advantage of this shape inference, but it is worth it even
    just for the ConvOp.  All that's necessary to do shape
    inference is 1) to mark shared inputs as having a particular
    shape, either via a .tag or some similar hacking; and 2) to
    add an optional In() argument to promise that inputs will
    have a certain shape (or even to have certain shapes in
    certain dimensions). We can't automatically infer the shape of
    shared variables as they can change of shape during the
    execution by default.  (NOT IMPLEMENTED YET, BUT IS IN TRAC)
    **Using Shape information in Optimizations**
    To use this shape information in OPTIMIZATIONS, use the
    ``shape_of`` dictionary.
    For example:
    .. code-block:: python
        try:
            shape_of = node.fgraph.shape_feature.shape_of
        except AttributeError:
            # This can happen when the mode doesn't include the ShapeFeature.
            return
        shape_of_output_zero = shape_of[node.output[0]]
    The ``shape_of_output_zero`` symbol will contain a tuple, whose
    elements are either integers or symbolic integers.
    TODO: check to see if the symbols are necessarily
    non-constant... or are integer literals sometimes Theano
    constants?? That would be confusing.
    """
    def get_node_infer_shape(self, node):
        try:
            shape_infer = node.op.infer_shape
        except AttributeError:
            shape_infer = self.default_infer_shape
        try:
            o_shapes = shape_infer(node,
                                   [self.shape_of[r] for r in node.inputs])
        except ShapeError:
            o_shapes = self.default_infer_shape(node, [self.shape_of[r] for
                                                       r in node.inputs])
        except NotImplementedError as e:
            raise NotImplementedError(
                'Code called by infer_shape failed raising a '
                'NotImplementedError. Raising NotImplementedError to '
                'indicate that a shape cannot be computed is no longer '
                'supported, and one should now use tensor.ShapeError '
                'instead. The original exception message is: %s' % e)
        except Exception as e:
            msg = ('Failed to infer_shape from Op %s.\nInput shapes: '
                   '%s\nException encountered during infer_shape: '
                   '%s\nException message: %s\nTraceback: %s') % (
                node.op, [self.shape_of[r] for r in node.inputs],
                type(e), str(e), traceback.format_exc())
            if config.on_shape_error == "raise":
                raise Exception(msg)
            else:
                _logger.warning(msg)
            o_shapes = self.default_infer_shape(
                node, [self.shape_of[r] for r in node.inputs])
        return o_shapes
    def get_shape(self, var, idx):
        """ Optimization can call this to get the current shape_i
        It is better to call this then use directly shape_of[var][idx]
        as this method should update shape_of if needed.
        TODO: Up to now, we don't update it in all cases. Update in all cases.
        """
        r = self.shape_of[var][idx]
        if (r.owner and
                isinstance(r.owner.op, Shape_i) and
                r.owner.inputs[0] not in var.fgraph.variables):
            assert var.owner
            node = var.owner
            # recur on inputs
            for i in node.inputs:
                if getattr(i, 'ndim', None) > 0:
                    self.get_shape(i, 0)
            o_shapes = self.get_node_infer_shape(node)
            assert len(o_shapes) == len(node.outputs)
            # Only change the variables and dimensions that would introduce
            # extra computation
            for new_shps, out in zip(o_shapes, node.outputs):
                if not hasattr(out, 'ndim'):
                    continue
                merged_shps = list(self.shape_of[out])
                changed = False
                for i in range(out.ndim):
                    n_r = merged_shps[i]
                    if (n_r.owner and
                            isinstance(n_r.owner.op, Shape_i) and
                            n_r.owner.inputs[0] not in var.fgraph.variables):
                        changed = True
                        merged_shps[i] = new_shps[i]
                if changed:
                    self.set_shape(out, merged_shps, override=True)
            r = self.shape_of[var][idx]
        return r
    def shape_ir(self, i, r):
        """Return symbolic r.shape[i] for tensor variable r, int i."""
        if hasattr(r.type, "broadcastable") and r.type.broadcastable[i]:
            return self.lscalar_one
        else:
            # Do not call make_node for test_value
            s = Shape_i(i)(r)
            try:
                s = get_scalar_constant_value(s)
            except NotScalarConstantError:
                pass
            return s
    def shape_tuple(self, r):
        """Return a tuple of symbolic shape vars for tensor variable r."""
        if not hasattr(r, 'ndim'):
            # This happen for NoneConst.
            return None
        return tuple([self.shape_ir(i, r) for i in xrange(r.ndim)])
    def default_infer_shape(self, node, i_shapes):
        """Return a list of shape tuple or None for the outputs of node.
        This function is used for Ops that don't implement infer_shape.
        Ops that do implement infer_shape should use the i_shapes parameter,
        but this default implementation ignores it.
        """
        rval = []
        for r in node.outputs:
            try:
                rval.append(self.shape_tuple(r))
            except AttributeError:
                rval.append(None)
        return rval
    def unpack(self, s_i, var):
        """Return a symbolic integer scalar for the shape element s_i.
        The s_i argument was produced by the infer_shape() of an Op subclass.
        var: the variable that correspond to s_i. This is just for
        error reporting.
        """
        # unpack the s_i that the Op returned
        assert s_i is not None
        if s_i == 1:
            # don't make the optimizer merge a zillion ones together
            # by always returning the same object to represent 1
            return self.lscalar_one
        if type(s_i) is float and int(s_i) == s_i:
            s_i = int(s_i)
        if (type(s_i) in integer_types or
                isinstance(s_i, numpy.integer) or
                (isinstance(s_i, numpy.ndarray) and s_i.ndim == 0)):
            # this shape is a constant
            if s_i < 0:
                msg = "There is a negative shape in the graph!"
                msg += gof.utils.get_variable_trace_string(var)
                raise ValueError(msg)
            return T.constant(s_i, dtype='int64')
        if type(s_i) in (tuple, list):
            # this dimension is the same as many of the inputs
            # which tells us that if one of the inputs is known,
            # the others all become known.
            # TODO: should be implemented in Elemwise, and Dot
            #
            # worst case, we loop over shape_of and replace things
            raise NotImplementedError(s_i)
        # s_i is x.shape[i] for some x, we change it to shape_of[x][i]
        if (s_i.owner and
                isinstance(s_i.owner.op, Subtensor) and
                s_i.owner.inputs[0].owner and
                isinstance(s_i.owner.inputs[0].owner.op, T.Shape)):
            assert s_i.ndim == 0
            assert len(s_i.owner.op.idx_list) == 1
            # The current Subtensor always put constant index in the graph.
            # This was not True in the past. So call the Subtensor function
            # that will return the right index.
            idx = get_idx_list(s_i.owner.inputs, s_i.owner.op.idx_list)
            assert len(idx) == 1
            idx = idx[0]
            try:
                i = get_scalar_constant_value(idx)
            except NotScalarConstantError:
                pass
            else:
                # Executed only if no exception was raised
                x = s_i.owner.inputs[0].owner.inputs[0]
                # x should already have been imported, and should be in shape_of.
                s_i = self.shape_of[x][i]
        if s_i.type.dtype in theano.tensor.integer_dtypes:
            if getattr(s_i.type, 'ndim', 0):
                raise TypeError('Shape element must be scalar', s_i)
            return s_i
        else:
            raise TypeError('Unsupported shape element',
                            s_i, type(s_i), getattr(s_i, 'type', None))
    def set_shape(self, r, s, override=False):
        """Assign the shape `s` to previously un-shaped variable `r`.
        Parameters
        ----------
        r : a variable
        s : None or a tuple of symbolic integers
        override : If False, it mean r is a new object in the fgraph.
            If True, it mean r is already in the fgraph and we want to
            override its shape.
        """
        if not override:
            assert r not in self.shape_of, 'r already in shape_of'
        if s is None:
            self.shape_of[r] = s
        else:
            if not isinstance(s, (tuple, list)):
                raise TypeError('shapes must be tuple/list', (r, s))
            if r.ndim != len(s):
                sio = StringIO()
                theano.printing.debugprint(r, file=sio, print_type=True)
                raise AssertionError(
                    "Something inferred a shape with %d dimensions "
                    "for a variable with %d dimensions"
                    " for the variable:\n%s" % (
                        len(s), r.ndim, sio.getvalue()))
            shape_vars = []
            for i in xrange(r.ndim):
                if (hasattr(r.type, 'broadcastable') and
                        r.type.broadcastable[i]):
                    shape_vars.append(self.lscalar_one)
                else:
                    shape_vars.append(self.unpack(s[i], r))
            assert all([not hasattr(r.type, "broadcastable") or
                        not r.type.broadcastable[i] or
                        # The two following comparison are a speed optimization
                        # But we never timed this speed optimization!
                        self.lscalar_one.equals(shape_vars[i]) or
                        self.lscalar_one.equals(
                            T.extract_constant(shape_vars[i]))
                        for i in xrange(r.ndim)])
            self.shape_of[r] = tuple(shape_vars)
            for sv in shape_vars:
                self.shape_of_reverse_index.setdefault(sv, set()).add(r)
    def update_shape(self, r, other_r):
        """Replace shape of r by shape of other_r.
        If, on some dimensions, the shape of other_r is not informative,
        keep the shape of r on those dimensions.
        """
        # other_r should already have a shape
        assert other_r in self.shape_of, ('other_r not in shape_of', other_r)
        other_shape = self.shape_of[other_r]
        # If other_shape has no information, call is pointless.
        if other_shape is None:
            return
        if r in self.shape_of:
            r_shape = self.shape_of[r]
        else:
            # If no info is known on r's shape, use other_shape
            self.set_shape(r, other_shape)
            return
        if (other_r.owner and r.owner and
                other_r.owner.inputs == r.owner.inputs and
                other_r.owner.op == r.owner.op):
            # We are doing a merge. So the 2 shapes graph will be the
            # same.  This is only a speed optimization to call
            # ancestors() less frequently.
            return
        # Merge other_shape with r_shape, giving the priority to other_shape
        merged_shape = []
        for i, ps in enumerate(other_shape):
            if r_shape is None and other_shape:
                merged_shape.append(other_shape[i])
            elif (ps.owner and
                    isinstance(getattr(ps.owner, 'op', None), Shape_i) and
                    ps.owner.op.i == i and
                    ps.owner.inputs[0] in (r, other_r)):
                # If other_shape[i] is uninformative, use r_shape[i].
                # For now, we consider 2 cases of uninformative other_shape[i]:
                #  - Shape_i(i)(other_r);
                #  - Shape_i(i)(r).
                merged_shape.append(r_shape[i])
            elif isinstance(r_shape[i], (Constant, integer_types)):
                # We do this to call less often ancestors and make
                # sure we have the simplest shape possible.
                merged_shape.append(r_shape[i])
            elif isinstance(other_shape[i], (Constant, integer_types)):
                # We do this to call less often ancestors and make
                # sure we have the simplest shape possible.
                merged_shape.append(other_shape[i])
            elif other_shape[i] == r_shape[i]:
                # This mean the shape is equivalent
                # We do not want to do the ancestor check in those cases
                merged_shape.append(r_shape[i])
            elif r_shape[i] in theano.gof.graph.ancestors([other_shape[i]]):
                # Another case where we want to use r_shape[i] is when
                # other_shape[i] actually depends on r_shape[i]. In that case,
                # we do not want to substitute an expression with another that
                # is strictly more complex. Such a substitution could also lead
                # to cycles: if (in the future) r_shape[i] gets replaced by an
                # expression of other_shape[i], other_shape[i] may end up
                # depending on itself.
                merged_shape.append(r_shape[i])
            else:
                merged_shape.append(other_shape[i])
        assert all([(not hasattr(r.type, "broadcastable") or
                     not r.type.broadcastable[i] and
                     not other_r.type.broadcastable[i]) or
                    # The two following comparison are a speed optimization
                    # But we never timed this speed optimization!
                    self.lscalar_one.equals(merged_shape[i]) or
                    self.lscalar_one.equals(
                        T.extract_constant(merged_shape[i], only_process_constants=True))
                    for i in xrange(r.ndim)])
        self.shape_of[r] = tuple(merged_shape)
        for sv in self.shape_of[r]:
            self.shape_of_reverse_index.setdefault(sv, set()).add(r)
    def set_shape_i(self, r, i, s_i):
        '''Replace element i of shape_of[r] by s_i'''
        assert r in self.shape_of
        prev_shape = self.shape_of[r]
        # prev_shape is a tuple, so we cannot change it inplace,
        # so we build another one.
        new_shape = []
        for j, s_j in enumerate(prev_shape):
            if j == i:
                new_shape.append(self.unpack(s_i, r))
            else:
                new_shape.append(s_j)
        assert all([not hasattr(r.type, "broadcastable") or
                    not r.type.broadcastable[idx] or
                    # The two following comparison are a speed optimization
                    # But we never timed this speed optimization!
                    self.lscalar_one.equals(new_shape[idx]) or
                    self.lscalar_one.equals(T.extract_constant(new_shape[idx]))
                    for idx in xrange(r.ndim)])
        self.shape_of[r] = tuple(new_shape)
        for sv in self.shape_of[r]:
            self.shape_of_reverse_index.setdefault(sv, set()).add(r)
    def init_r(self, r):
        '''Register r's shape in the shape_of dictionary.'''
        if r not in self.shape_of:
            try:
                self.set_shape(r, self.shape_tuple(r))
            except AttributeError:  # XXX: where would this come from?
                self.set_shape(r, None)
    def make_vector_shape(self, r):
        return make_vector(*self.shape_of[r])
    #
    # Feature interface
    #
    #
    def on_attach(self, fgraph):
        assert not hasattr(fgraph, 'shape_feature')
        fgraph.shape_feature = self
        # Must be local to the object as otherwise we reuse the same
        # variable for multiple fgraph!
        self.lscalar_one = T.constant(1, dtype='int64')
        assert self.lscalar_one.type == T.lscalar
        self.shape_of = {}
        # Variable -> tuple(scalars) or None  (All tensor vars map to tuple)
        self.scheduled = {}
        # Variable ->
        self.shape_of_reverse_index = {}
        # shape var -> graph v
        for node in fgraph.toposort():
            self.on_import(fgraph, node, reason='on_attach')
    def on_detach(self, fgraph):
        self.shape_of = {}
        self.scheduled = {}
        self.shape_of_reverse_index = {}
        del fgraph.shape_feature
    def on_import(self, fgraph, node, reason):
        if node.outputs[0] in self.shape_of:
            # this is a revert, not really an import
            for r in node.outputs + node.inputs:
                assert r in self.shape_of
            return
        for i, r in enumerate(node.inputs):
            # make sure we have shapes for the inputs
            self.init_r(r)
        o_shapes = self.get_node_infer_shape(node)
        # this is packed information
        # an element of o_shapes is either None or a tuple
        #   elements of the tuple can be either strings, or ints
        if len(o_shapes) != len(node.outputs):
            raise Exception(
                ('The infer_shape method for the Op "%s" returned a list ' +
                 'with the wrong number of element: len(o_shapes) = %d ' +
                 ' != len(node.outputs) = %d') % (str(node.op),
                                                  len(o_shapes),
                                                  len(node.outputs)))
        # Ensure shapes are in 'int64'. This is to make sure the assert
        # found in the `local_useless_subtensor` optimization does not fail.
        for sh_idx, sh in enumerate(o_shapes):
            if sh is None:
                continue
            if not isinstance(sh, (list, tuple)):
                raise ValueError("infer_shape of %s didn't return a list of"
                                 " list. It returned '%s'" % (str(node), str(o_shapes)))
            new_shape = []
            for i, d in enumerate(sh):
                # Note: we ignore any shape element that is not typed (i.e.,
                # does not have a 'dtype' attribute). This means there may
                # still remain int elements that are int32 on 32-bit platforms,
                # but this works with `local_useless_subtensor`, so for now we
                # keep it this way. See #266 for a better long-term fix.
                if getattr(d, 'dtype', 'int64') != 'int64':
                    assert d.dtype in theano.tensor.discrete_dtypes, (node, d.dtype)
                    assert str(d.dtype) != 'uint64', node
                    new_shape += sh[len(new_shape):i + 1]
                    if isinstance(d, T.Constant):
                        casted_d = T.constant(d.data, dtype='int64')
                    else:
                        casted_d = theano.tensor.cast(d, 'int64')
                    new_shape[i] = casted_d
            if new_shape:
                # We replace the shape with wrong dtype by the one with
                # 'int64'.
                new_shape += sh[len(new_shape):]
                o_shapes[sh_idx] = tuple(new_shape)
        for r, s in izip(node.outputs, o_shapes):
            self.set_shape(r, s)
    def on_change_input(self, fgraph, node, i, r, new_r, reason):
        if new_r not in self.shape_of:
            # It happen that the fgraph didn't called on_import for some
            # new_r.  This happen when new_r don't have an
            # owner(i.e. it is a constant or an input of the graph)
            # update_shape suppose that r and new_r are in shape_of.
            self.init_r(new_r)
        # This tells us that r and new_r must have the same shape if
        # we didn't know that the shapes are related, now we do.
        self.update_shape(new_r, r)
        # change_input happens in two cases:
        # 1) we are trying to get rid of r, or
        # 2) we are putting things back after a failed transaction.
        # In case 1, if r has a shape_i client, we will want to
        # replace the shape_i of r with the shape of new_r.  Say that
        # r is *scheduled*.
        # At that point, node is no longer a client of r, but of new_r
        for (shpnode, idx) in (r.clients + [(node, i)]):
            if isinstance(getattr(shpnode, 'op', None), Shape_i):
                idx = shpnode.op.i
                repl = self.shape_of[new_r][idx]
                if repl.owner is shpnode:
                    # This mean the replacement shape object is
                    # exactly the same as the current shape object. So
                    # no need for replacement. This happen for example
                    # with the InputToGpuOptimizer optimizer.
                    continue
                if (repl.owner and
                        repl.owner.inputs[0] is shpnode.inputs[0] and
                        isinstance(repl.owner.op, Shape_i) and
                        repl.owner.op.i == shpnode.op.i):
                    # The replacement is a shape_i of the same
                    # input. So no need to do this equivalent
                    # replacement.
                    continue
                if shpnode.outputs[0] in theano.gof.graph.ancestors([repl]):
                    raise InconsistencyError(
                        "This substitution would insert a cycle in the graph:"
                        "node: %s, i: %i, r: %s, new_r: %s"
                        % (node, i, r, new_r))
                self.scheduled[shpnode] = new_r
        # In case 2, if r is a variable that we've scheduled for shape update,
        # then we should cancel it.
        unscheduled = [k for k, v in self.scheduled.items() if v == r]
        for k in unscheduled:
            del self.scheduled[k]
        # In either case, r could be in shape_of.values(), that is, r itself
        # is the shape of  something. In that case, we want to update
        # the value in shape_of, to keep it up-to-date.
        for v in self.shape_of_reverse_index.get(r, []):
            # The reverse index is only approximate. It is not updated on
            # deletion of variables, or on change_input so it might be the
            # case that there are a few extra `v`'s in it that no longer have
            # a shape of r or possibly have been deleted from shape_of
            # entirely. The important thing is that it permits to recall
            # all variables with r in their shape.
            for ii, svi in enumerate(self.shape_of.get(v, [])):
                if svi == r:
                    self.set_shape_i(v, ii, new_r)
        self.shape_of_reverse_index[r] = set()
    def same_shape(self, x, y, dim_x=None, dim_y=None):
        """Return True if we are able to assert that x and y have the
        same shape.
        dim_x and dim_y are optional. If used, they should be an index
        to compare only 1 dimension of x and y.
        """
        sx = self.shape_of[x]
        sy = self.shape_of[y]
        if sx is None or sy is None:
            return False
        if dim_x is not None:
            sx = [sx[dim_x]]
        if dim_y is not None:
            sy = [sy[dim_y]]
        assert len(sx) == len(sy)
        # We look on each dimensions we want to compare.
        # If any of them can't be asserted to be equal, return False.
        # Otherwise, we return True at the end.
        for dx, dy in zip(sx, sy):
            if dx is dy:
                continue
            # Need to try to find that they are the same shape. We
            # need to compare the full graph. It could be slow. So I
            # just implement for now the case of Shape_i.
            if not dx.owner or not dy.owner:
                return False
            if (not isinstance(dx.owner.op, Shape_i) or
                    not isinstance(dy.owner.op, Shape_i)):
                return False
            opx = dx.owner.op
            opy = dy.owner.op
            if not (opx.i == opy.i):
                return False
            # FB I'm not sure if this handle correctly constants.
            if dx.owner.inputs[0] == dy.owner.inputs[0]:
                continue
            # To be sure to cover all case, call equal_computation.
            # Can't use theano.gof.graph.is_same_graph(dx, dy)
            # As it currently expect that dx and dy aren't in a FunctionGraph
            from theano.scan_module.scan_utils import equal_computations
            if not equal_computations([dx], [dy]):
                return False
        return True
class ShapeOptimizer(Optimizer):
    """Optimizer that serves to add ShapeFeature as an fgraph feature."""
    def add_requirements(self, fgraph):
        fgraph.attach_feature(ShapeFeature())
    def apply(self, fgraph):
        pass
class UnShapeOptimizer(Optimizer):
    """Optimizer remove ShapeFeature as an fgraph feature."""
    def apply(self, fgraph):
        for feature in fgraph._features:
            if isinstance(feature, ShapeFeature):
                fgraph.remove_feature(feature)
# Register it after merge1 optimization at 0. We don't want to track
# the shape of merged node.
theano.compile.mode.optdb.register('ShapeOpt', ShapeOptimizer(),
                                   0.1, 'fast_run', 'fast_compile')
# Not enabled by default for now. Some crossentropy opt use the
# shape_feature.  They are at step 2.01. uncanonicalize is at step
# 3. After it goes to 48.5 that move to the gpu. So 10 seem resonable.
theano.compile.mode.optdb.register('UnShapeOpt', UnShapeOptimizer(),
                                   10)
def local_elemwise_alloc_op(ElemwiseOP, AllocOP, DimShuffleOP):
    def local_elemwise_alloc(node):
        """
        elemwise(alloc(x, shp), ..., y.TensorType(BROADCAST CONDITION))
          -> elemwise(x, y.TensorType(BROADCAST CONDITION))
        elemwise(dimshuffle(alloc(x, shp)),... ,y.TensorType(BROADCAST CONDITION))
          -> elemwise(x.dimshuffle(...), y.TensorType(BROADCAST CONDITION))
        BROADCAST CONDITION: the condition is that the one input that are
        not to be optimized to have the same broadcast pattern as the
        output.
        We can change the alloc by a dimshuffle as the elemwise
        already have the shape info.  The dimshuffle will be faster
        to exec.
        """
        if not isinstance(node.op, ElemwiseOP):
            return False
        if len(node.outputs) > 1:
            # Ensure all outputs have the same broadcast pattern
            # This is a supposition that I'm not sure is always true.
            assert all([o.type.broadcastable ==
                        node.outputs[0].type.broadcastable for o in
                        node.outputs[1:]])
        # The broadcast pattern of the ouptut must match the broadcast
        # pattern of at least one of the inputs.
        if not any([i.type.broadcastable ==
                    node.outputs[0].type.broadcastable for i in node.inputs]):
            return False
        def dimshuffled_alloc(i):
            return (isinstance(i.owner.op, DimShuffleOP) and
                    i.owner.inputs[0].owner and
                    isinstance(i.owner.inputs[0].owner.op, AllocOP))
        # At least one input must have an owner that is either a AllocOP or a
        # DimShuffleOP with an owner that is a AllocOP -- otherwise there is
        # nothing to optimize.
        if not any([i.owner and (isinstance(i.owner.op, AllocOP) or
                                 dimshuffled_alloc(i)) for i in node.inputs]):
            return False
        # Search for input that we can use as a baseline for the dimensions.
        assert_op_idx = -1
        for idx, i in enumerate(node.inputs):
            if i.type.broadcastable == node.outputs[0].type.broadcastable:
                # Prefer an input that is not a AllocOP nor a DimShuffleOP of a
                # AllocOP so that all allocs can be optimized.
                if not (i.owner and (isinstance(i.owner.op, AllocOP) or
                        dimshuffled_alloc(i))):
                    assert_op_idx = idx
                    break
        # It may be the case that only AllocOP and DimShuffleOP of AllocOP exist.
        if assert_op_idx < 0:
            # We want to optimize as many allocs as possible. When
            # there is more than one then do all but one.  number of
            # inputs with alloc or dimshuffle alloc
            l2 = [i for i in node.inputs
                  if (i.owner and (isinstance(i.owner.op, AllocOP) or
                      dimshuffled_alloc(i)))]
            # If only 1 alloc or dimshuffle alloc, it is the one we
            # will use for the shape. So no alloc would be removed.
            if len(l2) > 1:
                # l containt inputs with alloc or dimshuffle alloc
                # only.  Its length will always be at least one, as we
                # checked that before
                l = [idx for idx, i in enumerate(node.inputs)
                     if i.broadcastable == node.outputs[0].broadcastable]
                assert_op_idx = l[0]  # The first one is as good as any to use.
            else:
                # Nothing would be optimized!
                return False
        assert_op = node.inputs[assert_op_idx]
        cmp_op = assert_op
        new_i = []
        same_shape = node.fgraph.shape_feature.same_shape
        for i in node.inputs:
            # Remove alloc
            if (i.owner and isinstance(i.owner.op, AllocOP) and
                    i.owner.inputs[0].type != i.owner.outputs[0].type):
                # when i.owner.inputs[0].type == i.owner.outputs[0].type we
                # will remove that alloc later
                assert i.type.ndim == cmp_op.ndim
                if theano.config.experimental.local_alloc_elemwise_assert:
                    get_shape = node.fgraph.shape_feature.get_shape
                    cond = []
                    for idx in xrange(i.type.ndim):
                        if (not i.type.broadcastable[idx] and
                                not same_shape(i, cmp_op, idx, idx)):
                            i_shp = get_shape(i, idx)
                            cmp_shp = get_shape(cmp_op, idx)
                            cond.append(T.eq(i_shp, cmp_shp))
                    if cond:
                        assert_op = assert_(assert_op, *cond)
                new_i.append(i.owner.inputs[0])
            # Remove Alloc in DimShuffle
            elif i.owner and dimshuffled_alloc(i):
                assert i.type.ndim == cmp_op.type.ndim
                if theano.config.experimental.local_alloc_elemwise_assert:
                    assert_cond = [T.eq(i.shape[idx], cmp_op.shape[idx])
                                   for idx in xrange(i.type.ndim)
                                   if not i.type.broadcastable[idx] and
                                   not same_shape(i, cmp_op, idx, idx)]
                    if assert_cond:
                        assert_op = assert_(assert_op, *assert_cond)
                alloc_input = i.owner.inputs[0].owner.inputs[0]
                if alloc_input.ndim != i.owner.inputs[0].ndim:
                    # The alloc can add dimension to the value
                    # We add a dimshuffle to add them.
                    # We let later optimization merge the multiple dimshuffle
                    nb_dim_to_add = i.owner.inputs[0].ndim - alloc_input.ndim
                    alloc_input = alloc_input.dimshuffle(
                        ['x'] * nb_dim_to_add +
                        list(range(alloc_input.ndim)))
                # We need to keep the dimshuffle. It could swap axes or
                # add dimensions anywhere.
                r_i = i.owner.op(alloc_input)
                # Copy stack trace from i to new_i
                copy_stack_trace(i, r_i)
                new_i.append(r_i)
            else:
                new_i.append(i)
        new_i[assert_op_idx] = assert_op
        ret = node.op(*new_i, return_list=True)
        # Copy over stack trace from previous outputs to new outputs.
        copy_stack_trace(node.outputs, ret)
        return ret
    return local_elemwise_alloc
# TODO, global optimizer that lift the assert to the beginning of the graph.
# TODO, optimize all inputs when possible -- currently when all inputs have
# an alloc all but one is optimized.
local_elemwise_alloc = register_specialize(
    gof.local_optimizer([T.Elemwise])(
        local_elemwise_alloc_op(T.Elemwise, T.Alloc, T.DimShuffle)),
    'local_alloc_elemwise')
@gof.local_optimizer([T.Elemwise])
def local_fill_sink(node):
    """
    f(fill(a, b), fill(c, d), e) -> fill(c, fill(a, f(b, d, e)))
    f need to be an elemwise that isn't a fill.
    """
    if (not hasattr(node, 'op') or
            not isinstance(node.op, T.Elemwise) or
            node.op == T.fill):
        return False
    models = []
    inputs = []
    for input in node.inputs:
        if input.owner and input.owner.op == T.fill:
            models.append(input.owner.inputs[0])
            inputs.append(input.owner.inputs[1])
        else:
            inputs.append(input)
    if not models:
        return False
    c = node.op(*inputs)
    for model in models:
        if model.type != c.type:
            c = T.fill(model, c)
    # The newly created node c doesn't has 'clients',
    # so this iteration is took place with node.outputs[0]
    replacements = {node.outputs[0]: c}
    for client, cl_idx in node.outputs[0].clients:
        if (hasattr(client, 'op') and
                isinstance(client.op, T.Elemwise) and
                not client.op == T.fill):
            client_inputs = client.inputs[:]
            client_inputs[cl_idx] = c
            new_client = client.op(*client_inputs)
            # Add clients to new_client
            new_client.owner.outputs[0].clients = client.outputs[0].clients
            r = local_fill_sink.transform(new_client.owner)
            if not r:
                continue
            replacements.update(r)
    return replacements
register_canonicalize(local_fill_sink)
@register_specialize
@register_stabilize
# @register_canonicalize  # We make full pass after the canonizer phase.
@gof.local_optimizer([T.fill])
def local_fill_to_alloc(node):
    """fill(s,v) -> alloc(v, shape(s))
    This is an important optimization because with the shape_to_shape_i
    optimization, the dependency on 's' is often removed.
    """
    if node.op == T.fill:
        r, v = node.inputs
        if v.type == node.outputs[0].type:
            # this is a useless fill, erase it.
            rval = [v]
        elif v.type.broadcastable == node.outputs[0].type.broadcastable:
            # this is a cast
            rval = [T.cast(v, node.outputs[0].type.dtype)]
        elif r.type.broadcastable == node.outputs[0].type.broadcastable:
            # we are broadcasting v somehow, but not r
            o = broadcast_like(v, r, node.fgraph, dtype=v.dtype)
            copy_stack_trace(node.outputs[0], o)
            rval = [o]
        else:
            # we are broadcasting both v and r,
            # the output shape must be computed
            #
            # TODO: implement this case (including a test!)
            #
            #  I think the strategy should be to extend the shorter
            #  shape vector with 1s (how?) and then take the
            #  elementwise max of the two.  - how to flag an error of
            #  shape mismatch where broadcasting should be illegal?
            return
            # TODO: cut out un-necessary dimshuffles of v
        assert rval[0].type == node.outputs[0].type, (
            'rval', rval[0].type, 'orig', node.outputs[0].type, 'node',
            node,)  # theano.printing.debugprint(node.outputs[0], file='str'))
        return rval
# Register this after stabilize at 1.5 to make sure stabilize don't
# get affected by less canonicalized graph due to alloc.
compile.optdb.register('local_fill_to_alloc',
                       in2out(local_fill_to_alloc),
                       1.51, 'fast_run')
# Needed to clean some extra alloc added by local_fill_to_alloc
compile.optdb.register('local_elemwise_alloc',
                       in2out(local_elemwise_alloc),
                       1.52, 'fast_run')
@register_canonicalize("fast_compile")
@register_useless
@gof.local_optimizer([T.fill])
def local_useless_fill(node):
    """fill(s,v) -> v
    This optimization is only needed in FAST_COMPILE to make the code
    more readable. Normally, it is done by the local_fill_to_alloc
    opt.
    """
    if node.op == T.fill:
        r, v = node.inputs
        if v.type == node.outputs[0].type:
            # this is a useless fill, erase it.
            # also, we don't need to copy over any stack traces here
            return [v]
@register_specialize
@register_stabilize
@register_canonicalize
@register_useless
@gof.local_optimizer([T.alloc])
def local_useless_alloc(node):
    """
    If the input type is the same as the output type (dtype and broadcast)
    there is no change in the shape of the input. So this is just a simple copy
    of the input. This is not needed.
    """
    op = node.op
    if not isinstance(op, Alloc):
        return False
    input = node.inputs[0]
    output = node.outputs[0]
    # Check if dtype and broadcast remain the same.
    if input.type == output.type:
        # We don't need to copy over any stack traces here
        return [input]
@register_specialize
@register_stabilize
@register_canonicalize
@gof.local_optimizer([T.alloc])
def local_canonicalize_alloc(node):
    """If the input type is the same as the output type (dtype and broadcast)
    there is no change in the shape of the input. So this is just a simple copy
    of the input. This is not needed. (as local_useless_alloc)
    Also, it will canonicalize alloc by creating Dimshuffle after the
    alloc to introduce the dimensions of constant size 1.
    See https://github.com/Theano/Theano/issues/4072 to know why this
    is needed.
    """
    op = node.op
    if not isinstance(op, Alloc):
        return False
    input = node.inputs[0]
    output = node.outputs[0]
    # Check if dtype and broadcast remain the same.
    if input.type == output.type:
        # We don't need to copy over any stack traces here
        return [input]
    # Allow local_merge_alloc to do its work first
    clients = getattr(output, 'clients', [])
    for client, i in clients:
        if client != "output" and isinstance(client.op, Alloc):
            return
    # Check if alloc adds a broadcastable dimension with shape 1.
    output_shape = node.inputs[1:]
    num_dims_with_size_1_added_to_left = 0
    for i in range(len(output_shape) - input.ndim):
        if extract_constant(output_shape[i], only_process_constants=True) == 1:
            num_dims_with_size_1_added_to_left += 1
        else:
            break
    new_output_shape = output_shape[num_dims_with_size_1_added_to_left:]
    if num_dims_with_size_1_added_to_left > 0 and len(new_output_shape) >= input.ndim:
        if output.broadcastable[num_dims_with_size_1_added_to_left:] == input.broadcastable:
            inner = input
        else:
            inner = op(*([input] + new_output_shape))
        dimshuffle_new_order = (['x'] * num_dims_with_size_1_added_to_left +
                                list(xrange(len(new_output_shape))))
        return [DimShuffle(inner.type.broadcastable, dimshuffle_new_order)(inner)]
# Don't register by default.
@gof.local_optimizer([T.AllocEmpty])
def local_alloc_empty_to_zeros(node):
    """This convert AllocEmpty to Alloc of 0.
    This help investigate NaN with NanGuardMode.  Not registered by
    default. To activate it, use the Theano flag
    optimizer_including=alloc_empty_to_zeros. This also enable
    the GPU version of this optimizations.
    """
    if isinstance(node.op, T.AllocEmpty):
        return [T.zeros(node.inputs, dtype=node.outputs[0].dtype)]
compile.optdb.register('local_alloc_empty_to_zeros',
                       in2out(local_alloc_empty_to_zeros),
                       # After move to gpu and merge2, before inplace.
                       49.3,
                       'alloc_empty_to_zeros',)
@register_specialize
@register_canonicalize
@gof.local_optimizer([T.Shape])
def local_shape_to_shape_i(node):
    if node.op == T.shape:
        # This optimization needs ShapeOpt and fgraph.shape_feature
        if not hasattr(node.fgraph, 'shape_feature'):
            return
        shape_feature = node.fgraph.shape_feature
        ret = shape_feature.make_vector_shape(node.inputs[0])
        # We need to copy over stack trace from input to output
        copy_stack_trace(node.outputs[0], ret)
        return [ret]
# TODO: Not sure what type of node we are expecting here
@register_specialize
@register_canonicalize
@gof.local_optimizer(None)
def local_track_shape_i(node):
    try:
        shape_feature = node.fgraph.shape_feature
    except AttributeError:
        return
    if node in shape_feature.scheduled:
        # Don't unschedule node as it could be reinserted in the
        # fgraph as we don't change it in the shapefeature internal
        # structure.
        assert isinstance(node.op, Shape_i)
        replacement = shape_feature.scheduled[node]
        return [shape_feature.shape_of[replacement][node.op.i]]
@register_specialize
@register_canonicalize
@gof.local_optimizer([Subtensor])
def local_subtensor_inc_subtensor(node):
    """
    Subtensor(SetSubtensor(x, y, idx), idx) -> y
    """
    if isinstance(node.op, Subtensor):
        x = node.inputs[0]
        if not x.owner or not isinstance(x.owner.op, IncSubtensor):
            return
        if not x.owner.op.set_instead_of_inc:
            return
        if (x.owner.inputs[2:] == node.inputs[1:] and
                tuple(x.owner.op.idx_list) == tuple(node.op.idx_list)):
            out = node.outputs[0]
            y = x.owner.inputs[1]
            # If the dtypes differ, cast y into x.dtype
            if x.dtype != y.dtype:
                y = y.astype(x.dtype)
            if out.type == y.type:
                # if x[idx] and y have the same type, directly return y
                return [y]
            else:
                # The difference is related to broadcasting pattern
                assert out.broadcastable != y.broadcastable
                # We have to alloc y to the shape of x[idx]
                x_subtensor = node.op(x.owner.inputs[0], *x.owner.inputs[2:])
                return [T.alloc(y, *x_subtensor.shape)]
        else:
            return
@register_specialize
@register_canonicalize
@gof.local_optimizer([Subtensor])
def local_subtensor_remove_broadcastable_index(node):
    """
    Remove broadcastable dimension with index 0 or -1
    a[:,:,:,0] -> a.dimshuffle(0,1,2), when
        a.broadcastable = (False, False, False, True)
    a[0,:,-1,:] -> a.dimshuffle(1,3), when
        a.broadcastable = (True, False, True, False)
    """
    if isinstance(node.op, Subtensor):
        idx = node.op.idx_list
    else:
        return
    remove_dim = []
    node_inputs_idx = 1
    for dim, elem in enumerate(idx):
        if isinstance(elem, (scalar.Scalar)):
            # The idx is a Scalar, ie a Type. This means the actual index
            # is contained in node.inputs[1]
            dim_index = node.inputs[node_inputs_idx]
            if type(dim_index) == theano.scalar.basic.ScalarConstant:
                dim_index = dim_index.value
            if dim_index in [0, -1] and node.inputs[0].broadcastable[dim]:
                remove_dim.append(dim)
                node_inputs_idx += 1
            else:
                return
        elif isinstance(elem, slice):
            if elem != slice(None):
                return
        elif isinstance(elem, (integer_types, numpy.integer)):
            if elem in [0, -1] and node.inputs[0].broadcastable[dim]:
                remove_dim.append(dim)
        else:
            raise TypeError('case not expected')
    if len(remove_dim) == 0:
        return
    else:
        all_dim = range(node.inputs[0].ndim)
        remain_dim = [x for x in all_dim if x not in remove_dim]
        return [node.inputs[0].dimshuffle(tuple(remain_dim))]
@register_specialize
@register_canonicalize('fast_compile_gpu')
@register_useless
@gof.local_optimizer([Subtensor, AdvancedSubtensor1])
def local_subtensor_make_vector(node):
    """
    Replace all subtensor(make_vector) like:
    [a,b,c][0] -> a
    [a,b,c][0:2] -> [a,b]
    Replace all AdvancedSubtensor1(make_vector) like:
    [a,b,c][[0,2]] -> [a,c]
    We can do this for constant indexes.
    """
    x = node.inputs[0]
    if not x.owner or x.owner.op != make_vector:
        return
    if isinstance(node.op, Subtensor):
        # This optimization needs ShapeOpt and fgraph.shape_feature
        try:
            idx, = node.op.idx_list
        except Exception:
            # 'how can you have multiple indexes into a shape?'
            raise
        if isinstance(idx, (scalar.Scalar, T.TensorType)):
            # The idx is a Scalar, ie a Type. This means the actual index
            # is contained in node.inputs[1]
            old_idx, idx = idx, node.inputs[1]
            assert idx.type == old_idx
    elif isinstance(node.op, AdvancedSubtensor1):
        idx = node.inputs[1]
    else:
        return
    if isinstance(idx, (integer_types, numpy.integer)):
        # We don't need to copy over any stack traces here
        return [x.owner.inputs[idx]]
    elif isinstance(idx, Variable):
        if idx.ndim == 0:
            # if it is a constant we can do something with it
            try:
                v = get_scalar_constant_value(idx, only_process_constants=True)
                if isinstance(v, numpy.integer):
                    # Python 2.4 wants to index only with Python integers
                    v = int(v)
                # We don't need to copy over any stack traces here
                try:
                    ret = [x.owner.inputs[v]]
                except IndexError:
                    raise NotScalarConstantError("Bad user graph!")
                return ret
            except NotScalarConstantError:
                pass
        elif idx.ndim == 1 and isinstance(idx, T.Constant):
            values = list(map(int, list(idx.value)))
            ret = make_vector(*[x.owner.inputs[v] for v in values])
            # Copy over stack trace from previous output to new output
            copy_stack_trace(node.outputs[0], ret)
            ret = T.patternbroadcast(ret, node.outputs[0].broadcastable)
            return [ret]
        else:
            raise TypeError('case not expected')
    elif isinstance(idx, slice):
        # it is a slice of ints and/or Variables
        # check subtensor to see if it can contain constant variables, and if
        # it can, then try to unpack them.
        try:
            const_slice = node.op.get_constant_idx(node.inputs,
                                                   allow_partial=False)[0]
            ret = make_vector(*x.owner.inputs[const_slice])
            # Copy over stack trace from previous outputs to new output
            copy_stack_trace(node.outputs, ret)
            ret = T.patternbroadcast(ret, node.outputs[0].broadcastable)
            return [ret]
        except NotScalarConstantError:
            pass
    else:
        raise TypeError('case not expected')
# TODO: the other optimization for and, or, xor, le and ge see ticket #496.
@register_useless
@register_canonicalize('fast_compile')
@register_specialize
@gof.local_optimizer([T.Elemwise])
def local_useless_elemwise(node):
    """
    eq(x,x) -> 1
    neq(x,x) -> 0
    mul(x) -> x
    add(x) -> x
    identity(x) -> x
    and(x,1) -> x
    and(x,0) -> zeros_like(x)
    or(x,0) -> x
    or(x,1) -> ones_like(x)
    xor(x,x) -> zeros_like(x)
    """
    if isinstance(node.op, T.Elemwise):
        # We call zeros_like and one_like with opt=True to generate a
        # cleaner graph.
        dtype = node.outputs[0].dtype
        if node.op.scalar_op == theano.scalar.eq and len(node.inputs) == 2:
            if node.inputs[0] == node.inputs[1]:
                # it is the same var in the graph. That will always be true
                ret = T.ones_like(node.inputs[0], dtype=dtype, opt=True)
                # Copy stack trace from input to constant output
                copy_stack_trace(node.outputs[0], ret)
                return [ret]
        elif node.op.scalar_op == theano.scalar.neq and len(node.inputs) == 2:
            if node.inputs[0] == node.inputs[1]:
                # it is the same var in the graph. That will always be false
                ret = T.zeros_like(node.inputs[0], dtype=dtype, opt=True)
                # Copy stack trace from input to constant output
                copy_stack_trace(node.outputs[0], ret)
                return [ret]
        elif node.op.scalar_op == theano.scalar.mul and len(node.inputs) == 1:
            # No need to copy over any stack trace
            return [node.inputs[0]]
        elif node.op.scalar_op == theano.scalar.add and len(node.inputs) == 1:
            # No need to copy over any stack trace
            return [node.inputs[0]]
        elif (node.op.scalar_op == theano.scalar.identity and
              len(node.inputs) == 1):
            return [node.inputs[0]]
        elif (isinstance(node.op.scalar_op, scalar.AND) and
              len(node.inputs) == 2):
            if isinstance(node.inputs[0], T.TensorConstant):
                const_val = T.extract_constant(node.inputs[0], only_process_constants=True)
                if not isinstance(const_val, Variable):
                    if const_val == 0:
                        return [T.zeros_like(node.inputs[1], dtype=dtype,
                                             opt=True)]
                    else:
                        return [node.inputs[1].astype(node.outputs[0].dtype)]
            if isinstance(node.inputs[1], T.TensorConstant):
                const_val = T.extract_constant(node.inputs[1], only_process_constants=True)
                if not isinstance(const_val, Variable):
                    if const_val == 0:
                        return [T.zeros_like(node.inputs[0], dtype=dtype,
                                             opt=True)]
                    else:
                        return [node.inputs[0].astype(node.outputs[0].dtype)]
        elif (isinstance(node.op.scalar_op, scalar.OR) and
              len(node.inputs) == 2):
            if isinstance(node.inputs[0], T.TensorConstant):
                const_val = T.extract_constant(node.inputs[0], only_process_constants=True)
                if not isinstance(const_val, Variable):
                    if const_val == 0:
                        return [node.inputs[1].astype(node.outputs[0].dtype)]
                    else:
                        return [T.ones_like(node.inputs[1], dtype=dtype,
                                            opt=True)]
            if isinstance(node.inputs[1], T.TensorConstant):
                const_val = T.extract_constant(node.inputs[1], only_process_constants=True)
                if not isinstance(const_val, Variable):
                    if const_val == 0:
                        return [node.inputs[0].astype(node.outputs[0].dtype)]
                    else:
                        return [T.ones_like(node.inputs[0], dtype=dtype,
                                            opt=True)]
        elif (isinstance(node.op.scalar_op, scalar.XOR) and
              len(node.inputs) == 2):
            if node.inputs[0] is node.inputs[1]:
                return [T.zeros_like(node.inputs[0], dtype=dtype, opt=True)]
@register_specialize
@gof.local_optimizer([T.Elemwise])
def local_alloc_unary(node):
    """unary(alloc(x, shp)) -> alloc(unary(x), shp)"""
    if isinstance(node.op, T.Elemwise) and len(node.inputs) == 1:
        a = node.inputs[0]
        if a.owner and isinstance(a.owner.op, T.Alloc):
            x = a.owner.inputs[0]
            shp = a.owner.inputs[1:]
            v = node.op(x)
            # T.alloc does not preserve the stacktrace of v,
            # so we need to copy it over from x.
            copy_stack_trace(node.outputs[0], v)
            ret = T.alloc(T.cast(v, node.outputs[0].dtype), *shp)
            # T.cast does not preserve the stacktrace of x,
            # so we need to copy it over to the output.
            copy_stack_trace([node.outputs[0], a], ret)
            return [ret]
@register_canonicalize
@register_specialize
@gof.local_optimizer([T.Elemwise])
def local_cast_cast(node):
    """cast(cast(x, dtype1), dtype2)
    when those contrain:
    dtype1 == dtype2
    TODO: the base dtype is the same (int, uint, float, complex)
          and the first cast cause an upcast.
    """
    if (not isinstance(node.op, T.Elemwise) or
            not isinstance(node.op.scalar_op, scalar.Cast)):
        return
    x = node.inputs[0]
    if (not x.owner or
            not isinstance(x.owner.op, T.Elemwise) or
            not isinstance(x.owner.op.scalar_op, scalar.Cast)):
        return
    if node.op.scalar_op.o_type == x.owner.op.scalar_op.o_type:
        # We don't need to copy over any stack traces here
        return [x]
@register_canonicalize
@register_specialize
@gof.local_optimizer([T.Elemwise])
def local_func_inv(node):
    """
    Check for two consecutive operations that are functional inverses
    and remove them from the function graph.
    """
    inv_pairs = (
        (basic.Deg2Rad, basic.Rad2Deg),
        (basic.Cosh, basic.ArcCosh),
        (basic.Tanh, basic.ArcTanh),
        (basic.Sinh, basic.ArcSinh),
        (basic.Conj, basic.Conj),
        (basic.Neg, basic.Neg),
        (basic.Inv, basic.Inv),
    )
    x = node.inputs[0]
    if not isinstance(node.op, T.Elemwise):
        return
    if (not x.owner or not isinstance(x.owner.op, T.Elemwise)):
        return
    prev_op = x.owner.op.scalar_op
    node_op = node.op.scalar_op
    for inv_pair in inv_pairs:
        if is_inverse_pair(node_op, prev_op, inv_pair):
            # We don't need to copy stack trace, because the optimization
            # is trivial and maintains the earlier stack trace
            return x.owner.inputs
    return
def is_inverse_pair(node_op, prev_op, inv_pair):
    """
    Given two consecutive operations, check if they are the
    provided pair of inverse functions.
    """
    node_is_op0 = isinstance(node_op, inv_pair[0])
    node_is_op1 = isinstance(node_op, inv_pair[1])
    prev_is_op0 = isinstance(prev_op, inv_pair[0])
    prev_is_op1 = isinstance(prev_op, inv_pair[1])
    return (node_is_op0 and prev_is_op1) or (node_is_op1 and prev_is_op0)
class Assert(T.Op):
    """
    Implements assertion in a computational graph.
    Returns the first parameter if the condition is true, otherwise, triggers
    AssertionError.
    Notes
    -----
    This Op is a debugging feature. It can be removed from the graph
    because of optimizations, and can hide some possible optimizations to
    the optimizer. Specifically, removing happens if it can be determined
    that condition will always be true. Also, the output of the Op must be
    used in the function computing the graph, but it doesn't have to be
    returned.
    Examples
    --------
    >>> import theano
    >>> T = theano.tensor
    >>> x = T.vector('x')
    >>> assert_op = T.opt.Assert()
    >>> func = theano.function([x], assert_op(x, x.size<2))
    """
    _f16_ok = True
    __props__ = ('msg',)
    view_map = {0: [0]}
    check_input = False
    def __init__(self, msg="Theano Assert failed!"):
        self.msg = msg
    def __setstate__(self, attrs):
        self.__dict__.update(attrs)
        if not hasattr(self, 'msg'):
            self.msg = "Theano Assert failed!"
    def make_node(self, value, *conds):
        if not isinstance(value, Variable):
            value = T.as_tensor_variable(value)
        cond = [T.as_tensor_variable(c) for c in conds]
        assert numpy.all([c.type.ndim == 0 for c in cond])
        return gof.Apply(self, [value] + cond, [value.type()])
    def perform(self, node, inputs, out_):
        out, = out_
        v = inputs[0]
        out[0] = v
        assert numpy.all(inputs[1:]), self.msg
    def grad(self, input, output_gradients):
        return output_gradients + [DisconnectedType()()] * (len(input) - 1)
    def connection_pattern(self, node):
        return [[1]] + [[0]] * (len(node.inputs) - 1)
    def c_code(self, node, name, inames, onames, sub):
        value = inames[0]
        out = onames[0]
        check = []
        fail = sub['fail']
        msg = self.msg.replace('"', '\\"').replace('\n', '\\n')
        for idx in xrange(len(inames) - 1):
            i = inames[idx + 1]
            dtype = node.inputs[idx + 1].dtype
            check.append('if(!((npy_%(dtype)s*)PyArray_DATA(%(i)s))[0])'
                         '{PyErr_SetString(PyExc_AssertionError,"%(msg)s");'
                         '%(fail)s}' % locals())
        check = "\n".join(check)
        return """
        %(check)s
        Py_XDECREF(%(out)s);
        %(out)s = %(value)s;
        Py_INCREF(%(value)s);
        """ % locals()
    def c_code_cache_version(self):
        return (3, 0)
    def infer_shape(self, node, input_shapes):
        return [input_shapes[0]]
assert_ = Assert()
# Unittest.assert_ is a deprecated name for assertTrue.
# 2to3 convert theano.tensor.opt.assert_ to theano.tensor.opt.assertTrue
# So I define a new name as a work around.
assert_op = assert_
@register_specialize
@gof.local_optimizer([Assert])
def local_remove_useless_assert(node):
    if isinstance(node.op, Assert):
        cond = []
        for c in node.inputs[1:]:
            try:
                const = get_scalar_constant_value(c)
                if 0 != const.ndim or const == 0:
                    # Should we raise an error here? How to be sure it
                    # is not catched?
                    cond.append(c)
            except NotScalarConstantError:
                cond.append(c)
        if len(cond) == 0:
            # We don't need to copy over any stack traces here
            return [node.inputs[0]]
        if len(cond) != len(node.inputs) - 1:
            ret = assert_(node.inputs[0], *cond)
            # We copy over stack trace from the output of the original assert
            copy_stack_trace(node.outputs[0], ret)
            return [ret]
@gof.local_optimizer([Assert])
def local_remove_all_assert(node):
    """An optimization disabled by default that removes all asserts from
    the graph.
    Notes
    -----
    See the :ref:`unsafe` section to know how to enable it.
    """
    if not isinstance(node.op, Assert):
        return
    # We don't need to copy over any stack traces here
    return [node.inputs[0]]
# Disabled by default
compile.optdb['canonicalize'].register('local_remove_all_assert',
                                       local_remove_all_assert,
                                       'unsafe',
                                       use_db_name_as_tag=False)
compile.optdb['stabilize'].register('local_remove_all_assert',
                                    local_remove_all_assert,
                                    'unsafe',
                                    use_db_name_as_tag=False)
compile.optdb['specialize'].register('local_remove_all_assert',
                                     local_remove_all_assert,
                                     'unsafe',
                                     use_db_name_as_tag=False)
compile.optdb['useless'].register('local_remove_all_assert',
                                  local_remove_all_assert,
                                  'unsafe',
                                  use_db_name_as_tag=False)
#######################
# Constant Canonicalization
############################
@register_canonicalize
@gof.local_optimizer([T.Elemwise])
def local_upcast_elemwise_constant_inputs(node):
    """This explicitly upcasts constant inputs to elemwise Ops, when
    those Ops do implicit upcasting anyway.
    Rationale: it helps merge things like (1-x) and (1.0 - x).
    """
    if len(node.outputs) > 1:
        return
    try:
        shape_i = node.fgraph.shape_feature.shape_i
    except AttributeError:
        shape_i = None
    if isinstance(node.op, T.Elemwise):
        scalar_op = node.op.scalar_op
        # print "aa", scalar_op.output_types_preference
        if (getattr(scalar_op, 'output_types_preference', None)
                in (T.scal.upgrade_to_float, T.scal.upcast_out)):
            # this is the kind of op that we can screw with the input
            # dtypes by upcasting explicitly
            output_dtype = node.outputs[0].type.dtype
            new_inputs = []
            for i in node.inputs:
                if i.type.dtype == output_dtype:
                    new_inputs.append(i)
                else:
                    try:
                        # works only for scalars
                        cval_i = get_scalar_constant_value(i,
                                                           only_process_constants=True)
                        if all(i.broadcastable):
                            new_inputs.append(T.shape_padleft(
                                T.cast(cval_i, output_dtype),
                                i.ndim))
                        else:
                            if shape_i is None:
                                return
                            new_inputs.append(
                                T.alloc(T.cast(cval_i, output_dtype),
                                        *[shape_i(d)(i)
                                          for d in xrange(i.ndim)]))
                            # print >> sys.stderr, "AAA",
                            # *[Shape_i(d)(i) for d in xrange(i.ndim)]
                    except NotScalarConstantError:
                        # for the case of a non-scalar
                        if isinstance(i, T.TensorConstant):
                            new_inputs.append(T.cast(i, output_dtype))
                        else:
                            new_inputs.append(i)
            if new_inputs != node.inputs:
                rval = [node.op(*new_inputs)]
                if rval[0].type != node.outputs[0].type:
                    # This can happen for example when floatX=float32
                    # and we do the true division between and int64
                    # and a constant that will get typed as int8.
                    # As this is just to allow merging more case, if
                    # the upcast don't work, we can just skip it.
                    return
                # Copy over output stacktrace from before upcasting
                copy_stack_trace(node.outputs[0], rval)
                return rval
##################
# Subtensor opts #
##################
@register_useless
@register_canonicalize
@register_specialize
@gof.local_optimizer([IncSubtensor])
def local_useless_inc_subtensor(node):
    """
    Remove IncSubtensor, when we overwrite the full inputs with the
    new value.
    """
    if not isinstance(node.op, IncSubtensor):
        return
    if node.op.set_instead_of_inc is False:
        # This is an IncSubtensor, so the init value must be zeros
        try:
            c = get_scalar_constant_value(node.inputs[0],
                                          only_process_constants=True)
            if c != 0:
                return
        except NotScalarConstantError:
            return
    if (node.inputs[0].ndim != node.inputs[1].ndim or
            node.inputs[0].broadcastable != node.inputs[1].broadcastable):
        # FB: I didn't check if this case can happen, but this opt
        # don't support it.
        return
    # We have a SetSubtensor or an IncSubtensor on zeros
    # If is this IncSubtensor useful?
    # Check that we keep all the original data.
    # Put the constant inputs in the slice.
    idx_cst = get_idx_list(node.inputs[1:], node.op.idx_list)
    if all(isinstance(e, slice) and e.start is None and
           e.stop is None and (e.step is None or T.extract_constant(e.step,
                               only_process_constants=True) == -1)
           for e in idx_cst):
        # IncSubtensor broadcast node.inputs[1] on node.inputs[0]
        # based on run time shapes, so we must check they are the same.
        if not hasattr(node.fgraph, 'shape_feature'):
            return
        if not node.fgraph.shape_feature.same_shape(node.inputs[0],
                                                    node.inputs[1]):
            return
        # There is no reverse, so we don't need a replacement.
        if all(e.step is None
               for e in node.op.idx_list):
            # They are the same shape, so we can remore this IncSubtensor
            return [node.inputs[1]]
        ret = Subtensor(node.op.idx_list)(*node.inputs[1:])
        # Copy over previous output stacktrace
        copy_stack_trace(node.outputs, ret)
        return [ret]
@register_canonicalize
@gof.local_optimizer([AdvancedIncSubtensor1])
def local_set_to_inc_subtensor(node):
    """
    AdvancedIncSubtensor1(x, x[ilist]+other, ilist, set_instead_of_inc=True) ->
    AdvancedIncSubtensor1(x, other, ilist, set_instead_of_inc=False)
    """
    if (isinstance(node.op, AdvancedIncSubtensor1) and
            node.op.set_instead_of_inc and
            node.inputs[1].owner and
            isinstance(node.inputs[1].owner.op, Elemwise) and
            isinstance(node.inputs[1].owner.op.scalar_op, scalar.Add)):
        addn = node.inputs[1].owner
        subn = None
        other = None
        if (addn.inputs[0].owner and
                isinstance(addn.inputs[0].owner.op, AdvancedSubtensor1)):
            subn = addn.inputs[0].owner
            other = addn.inputs[1]
        elif (addn.inputs[1].owner and
              isinstance(addn.inputs[1].owner.op, AdvancedSubtensor1)):
            subn = addn.inputs[1].owner
            other = addn.inputs[0]
        else:
            return
        if (subn.inputs[1] != node.inputs[2] or
                subn.inputs[0] != node.inputs[0]):
            return
        ret = advanced_inc_subtensor1(node.inputs[0], other, node.inputs[2])
        # Copy over previous output stacktrace
        # Julian: I'm not sure about this at all...
        copy_stack_trace(node.outputs, ret)
        return [ret]
@register_useless
@register_canonicalize
@register_specialize
@gof.local_optimizer([Subtensor])
def local_useless_slice(node):
    """
    Remove Subtensor of the form X[0, :] -> X[0]
    """
    if isinstance(node.op, Subtensor):
        slices = get_idx_list(node.inputs, node.op.idx_list)
        last_slice = len(slices)
        for s in slices[::-1]:
            # check if slice and then check slice indices
            if (isinstance(s, slice) and s.start is None and s.stop is None and
                    (s.step is None or T.extract_constant(s.step,
                                                          only_process_constants=True) == 1)):
                last_slice -= 1
            else:
                break
        # check if we removed something
        if last_slice < len(slices):
            subtens = Subtensor(slices[:last_slice])
            sl_ins = Subtensor.collapse(slices[:last_slice],
                                        lambda x: isinstance(x, T.Variable))
            out = subtens(node.inputs[0], *sl_ins)
            # Copy over previous output stacktrace
            copy_stack_trace(node.outputs, out)
            return [out]
@register_canonicalize
@register_specialize
@gof.local_optimizer([Subtensor, AdvancedSubtensor1])
def local_useless_subtensor(node):
    """
    Remove Subtensor/AdvancedSubtensor1 if it takes the full input. In the
    AdvancedSubtensor1 case, the full input is taken when the indices are
    equivalent to `arange(0, input.shape[0], 1)` using either an explicit
    list/vector or the ARange op.
    """
    # If the optimization is tried over a node that is not a part of graph before
    if not hasattr(node, 'fgraph'):
        return
    # This optimization needs ShapeOpt and fgraph.shape_feature
    if not hasattr(node.fgraph, 'shape_feature'):
        return
    shape_of = node.fgraph.shape_feature.shape_of
    if isinstance(node.op, Subtensor):
        cdata = node.op.get_constant_idx(node.inputs, allow_partial=True,
                                         only_process_constants=True)


Answer:         for pos, idx in enumerate(cdata):