""" Tensor optimizations addressing the ops in math.py."""

import itertools
import operator
from functools import partial, reduce

import numpy as np

import aesara.scalar.basic as aes
import aesara.scalar.math as aes_math
from aesara.graph.basic import Constant, Variable
from aesara.graph.opt import (
    LocalOptGroup,
    LocalOptimizer,
    PatternSub,
    copy_stack_trace,
    in2out,
    local_optimizer,
)
from aesara.graph.opt_utils import get_clients_at_depth
from aesara.misc.safe_asarray import _asarray
from aesara.raise_op import assert_op
from aesara.tensor.basic import (
    Alloc,
    Join,
    MakeVector,
    alloc,
    as_tensor_variable,
    cast,
    constant,
    extract_constant,
    fill,
    get_scalar_constant_value,
    ones_like,
    switch,
    zeros_like,
)
from aesara.tensor.basic_opt import (
    FusionOptimizer,
    broadcast_like,
    encompasses_broadcastable,
    fuse_seqopt,
    local_fill_sink,
    register_canonicalize,
    register_specialize,
    register_specialize_device,
    register_stabilize,
    register_uncanonicalize,
    register_useless,
)
from aesara.tensor.elemwise import CAReduce, DimShuffle, Elemwise
from aesara.tensor.exceptions import NotScalarConstantError
from aesara.tensor.math import (
    All,
    Any,
    Dot,
    NonZeroCAReduce,
    Prod,
    ProdWithoutZeros,
    Sum,
)
from aesara.tensor.math import abs as at_abs
from aesara.tensor.math import (
    add,
    dot,
    eq,
    erf,
    erfc,
    exp,
    expm1,
    ge,
    int_div,
    isinf,
    le,
    log,
    log1mexp,
    log1p,
    makeKeepDims,
)
from aesara.tensor.math import max as at_max
from aesara.tensor.math import maximum, mul, neg
from aesara.tensor.math import pow as at_pow
from aesara.tensor.math import prod, reciprocal, sgn, sigmoid, softplus, sqr, sqrt, sub
from aesara.tensor.math import sum as at_sum
from aesara.tensor.math import true_div
from aesara.tensor.shape import Shape, Shape_i
from aesara.tensor.subtensor import Subtensor
from aesara.tensor.type import (
    uint_dtypes,
    values_eq_approx_remove_inf,
    values_eq_approx_remove_inf_nan,
    values_eq_approx_remove_nan,
)
from aesara.tensor.var import TensorConstant, get_unique_value


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 get_constant(v):
    """

    Returns
    -------
    object
        A numeric constant if v is a Constant or, well, a
        numeric constant. If v is a plain Variable, returns None.

    """
    if isinstance(v, Constant):
        unique_value = get_unique_value(v)
        if unique_value is not None:
            data = unique_value
        else:
            data = v.data
        if data.ndim == 0:
            return data
        else:
            return None
    elif isinstance(v, Variable):
        return None
    else:
        return v


def fill_chain(new_out, orig_inputs):
    for i in orig_inputs:
        new_out = fill(i, new_out)
    return [new_out]


@register_canonicalize
@register_stabilize
@local_optimizer([Dot])
def local_0_dot_x(fgraph, node):
    if not isinstance(node.op, 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 = constant(0, dtype=node.outputs[0].type.dtype)
        if x.ndim == 2 and y.ndim == 2:
            constant_zero = assert_op(constant_zero, eq(x.shape[1], y.shape[0]))
            return [alloc(constant_zero, x.shape[0], y.shape[1])]
        elif x.ndim == 1 and y.ndim == 2:
            constant_zero = assert_op(constant_zero, eq(x.shape[0], y.shape[0]))
            return [alloc(constant_zero, y.shape[1])]
        elif x.ndim == 2 and y.ndim == 1:
            constant_zero = assert_op(constant_zero, eq(x.shape[1], y.shape[0]))
            return [alloc(constant_zero, x.shape[0])]
        elif x.ndim == 1 and y.ndim == 1:
            constant_zero = assert_op(constant_zero, eq(x.shape[0], y.shape[0]))
            return [constant_zero]


@register_canonicalize
@local_optimizer([DimShuffle])
def local_lift_transpose_through_dot(fgraph, node):
    """Perform the rewrite ``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, DimShuffle) and node.op.new_order == (1, 0)):
        return False
    if not (node.inputs[0].owner and isinstance(node.inputs[0].owner.op, 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 = [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


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)


@register_canonicalize
@register_specialize
@local_optimizer([Elemwise])
def local_func_inv(fgraph, node):
    """
    Check for two consecutive operations that are functional inverses
    and remove them from the function graph.

    """
    inv_pairs = (
        (aes.Deg2Rad, aes.Rad2Deg),
        (aes.Cosh, aes.ArcCosh),
        (aes.Tanh, aes.ArcTanh),
        (aes.Sinh, aes.ArcSinh),
        (aes.Conj, aes.Conj),
        (aes.Neg, aes.Neg),
        (aes.Reciprocal, aes.Reciprocal),
    )
    x = node.inputs[0]

    if not isinstance(node.op, Elemwise):
        return
    if not x.owner or not isinstance(x.owner.op, 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
            ottype = node.out.dtype
            inp = x.owner.inputs[0]
            # Functions may have casted integer input to float
            if inp.dtype != ottype:
                inp = cast(inp, ottype)
            return [inp]

    return


@register_canonicalize
@register_specialize
@local_optimizer([Elemwise])
def local_exp_log(fgraph, node):
    x = node.inputs[0]

    if not isinstance(node.op, Elemwise):
        return
    if not x.owner or not isinstance(x.owner.op, Elemwise):
        return

    prev_op = x.owner.op.scalar_op
    node_op = node.op.scalar_op

    # Case for log(exp(x)) -> x
    if isinstance(prev_op, aes.Exp) and isinstance(node_op, aes.Log):
        new_out = x.owner.inputs[0]
        old_out = node.outputs[0]
        # Exp may have cast integer input to float
        if new_out.dtype != old_out.dtype:
            new_out = cast(new_out, old_out.dtype)
        return [new_out]

    # Case for log1p(expm1(x)) -> x
    if isinstance(prev_op, aes.Expm1) and isinstance(node_op, aes.Log1p):
        new_out = x.owner.inputs[0]
        old_out = node.outputs[0]
        # Expm1 may have cast integer input to float
        if new_out.dtype != old_out.dtype:
            new_out = cast(new_out, old_out.dtype)
        return [new_out]

    # Case for exp(softplus(x)) aka exp(log1pexp) -> 1 + exp(x)
    if isinstance(prev_op, aes_math.Softplus) and isinstance(node_op, aes.Exp):
        x = x.owner.inputs[0]
        return [add(1, exp(x))]

    # Case for expm1(softplus(x)) aka expm1(log1pexp) -> exp(x)
    if isinstance(prev_op, aes_math.Softplus) and isinstance(node_op, aes.Expm1):
        x = x.owner.inputs[0]
        return [exp(x)]


@register_specialize
@local_optimizer([Elemwise])
def local_exp_log_nan_switch(fgraph, node):
    # Rewrites of the kind exp(log...(x)) that require a `nan` switch
    x = node.inputs[0]

    if not isinstance(node.op, Elemwise):
        return
    if not x.owner or not isinstance(x.owner.op, Elemwise):
        return

    prev_op = x.owner.op.scalar_op
    node_op = node.op.scalar_op

    # Case for exp(log(x)) -> x
    if isinstance(prev_op, aes.Log) and isinstance(node_op, aes.Exp):
        x = x.owner.inputs[0]
        old_out = node.outputs[0]
        new_out = switch(ge(x, 0), x, np.asarray(np.nan, old_out.dtype))
        return [new_out]

    # Case for exp(log1p(x)) -> x + 1
    if isinstance(prev_op, aes.Log1p) and isinstance(node_op, aes.Exp):
        x = x.owner.inputs[0]
        old_out = node.outputs[0]
        new_out = switch(ge(x, -1), add(1, x), np.asarray(np.nan, old_out.dtype))
        return [new_out]

    # Case for expm1(log(x)) -> x - 1
    if isinstance(prev_op, aes.Log) and isinstance(node_op, aes.Expm1):
        x = x.owner.inputs[0]
        old_out = node.outputs[0]
        new_out = switch(ge(x, 0), sub(x, 1), np.asarray(np.nan, old_out.dtype))
        return [new_out]

    # Case for expm1(log1p(x)) -> x
    if isinstance(prev_op, aes.Log1p) and isinstance(node_op, aes.Expm1):
        x = x.owner.inputs[0]
        old_out = node.outputs[0]
        new_out = switch(ge(x, -1), x, np.asarray(np.nan, old_out.dtype))
        return [new_out]

    # Case for exp(log1mexp(x)) -> 1 - exp(x)
    if isinstance(prev_op, aes_math.Log1mexp) and isinstance(node_op, aes.Exp):
        x = x.owner.inputs[0]
        old_out = node.outputs[0]
        new_out = switch(le(x, 0), sub(1, exp(x)), np.asarray(np.nan, old_out.dtype))
        return [new_out]

    # Case for expm1(log1mexp(x)) -> -exp(x)
    if isinstance(prev_op, aes_math.Log1mexp) and isinstance(node_op, aes.Expm1):
        x = x.owner.inputs[0]
        old_out = node.outputs[0]
        new_out = switch(le(x, 0), neg(exp(x)), np.asarray(np.nan, old_out.dtype))
        return [new_out]


@register_canonicalize
@register_specialize
@local_optimizer([Sum])
def local_sumsqr2dot(fgraph, node):
    """
    This optimization detects
    ``at.sqr(W.dimshuffle("x", 0, 1) * G.dimshuffle(0, "x", 1) ).sum(axis=(1, 2))``
    and converts it to ``at.dot(at.sqr(G), at.sqr(W).sum(axis=0))``.
    """
    if (
        isinstance(node.op, Sum)
        and isinstance(node.op.scalar_op, aes.Add)
        and node.op.axis == (1, 2)
    ):
        in1 = node.inputs[0]
        out = node.outputs[0]

        if (
            in1.owner
            and isinstance(in1.owner.op, Elemwise)
            and isinstance(in1.owner.op.scalar_op, aes.Sqr)
        ):
            in_sqr = in1.owner.inputs[0]
            if (
                in_sqr.owner
                and isinstance(in_sqr.owner.op, Elemwise)
                and isinstance(in_sqr.owner.op.scalar_op, aes.Mul)
                and len(in_sqr.owner.inputs) == 2
            ):
                in_mul1, in_mul2 = in_sqr.owner.inputs

                if (
                    isinstance(in_mul1.owner.op, DimShuffle)
                    and in_mul1.owner.op.new_order == ("x", 0, 1)
                    and isinstance(in_mul2.owner.op, DimShuffle)
                    and in_mul2.owner.op.new_order == (0, "x", 1)
                ):
                    W = in_mul1.owner.inputs[0]
                    G = in_mul2.owner.inputs[0]

                    new_out = dot(sqr(G), sqr(W).sum(axis=0))
                    if new_out.dtype != out.dtype:
                        new_out = cast(new_out, dtype=out.dtype)
                    return [new_out]


@register_stabilize
@register_specialize
@register_canonicalize
@local_optimizer([Elemwise])
def local_expm1(fgraph, node):
    """
    This optimization detects exp(a)-1 and converts this to expm1(a).
    """
    if isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, aes.Sub):
        in1, in2 = node.inputs
        out = node.outputs[0]

        if (
            in1.owner
            and isinstance(in1.owner.op, Elemwise)
            and isinstance(in1.owner.op.scalar_op, aes.Exp)
            and extract_constant(in2, only_process_constants=False) == 1
        ):
            in11 = in1.owner.inputs[0]
            new_out = expm1(in11)

            if new_out.dtype != out.dtype:
                new_out = cast(new_out, dtype=out.dtype)

            if not out.type.is_super(new_out.type):
                return
            return [new_out]


@register_specialize
@register_canonicalize
@local_optimizer([mul])
def local_mul_switch_sink(fgraph, node):
    """
    This optimization makes the following changes in the graph:
    ``at.mul(A, at.switch(cond, 0, iff), B)`` -> ``at.switch(cond, 0, at.mul(A, B, iff))``
    ``at.mul(A, at.switch(cond, ift, 0), B)`` -> ``at.switch(cond, at.mul(A, B, ift), 0)``
    ``A`` and ``B`` being several (or none) symbolic variables.
    This is useful because ``A`` and ``B`` may not be numerically stable and give
    NaN or inf values for cases where the switch returns 0.
    With this optimization ``at.grad(at.switch(...))`` has the right behavior.

    Examples
    --------

        x -> f(x)
        x -> g(x)
        y = at.switch(cond, f(x), g(x))

    without the optimization:

        at.grad(y, x) -> grad(f(x), x) * grad(y, f(x)) + grad(g(x), x) * grad(y, g(x))

    with the optimization

        at.grad(y, x) -> switch(cond, grad(f(x), x), 0) + switch(cond, 0, grad(g(x), x))

    This will be particularly useful for the lazy ``if`` because we skip an entire
    part of the graph.

    """
    if node.op != mul:
        return False
    for idx, i in enumerate(node.inputs):
        if i.owner and i.owner.op == switch:
            switch_node = i.owner
            try:
                if (
                    get_scalar_constant_value(
                        switch_node.inputs[1], only_process_constants=True
                    )
                    == 0.0
                ):
                    listmul = node.inputs[:idx] + node.inputs[idx + 1 :]
                    fmul = mul(*(listmul + [switch_node.inputs[2]]))

                    # Copy over stacktrace for elementwise multiplication op
                    # from previous elementwise multiplication op.
                    # An error in the multiplication (e.g. errors due to
                    # inconsistent shapes), will point to the
                    # multiplication op.
                    copy_stack_trace(node.outputs, fmul)

                    fct = [switch(switch_node.inputs[0], 0, fmul)]
                    fct[0].tag.values_eq_approx = values_eq_approx_remove_nan

                    # Copy over stacktrace for switch op from both previous
                    #  elementwise multiplication op and previous switch op,
                    # because an error in this part can be caused by either
                    # of the two previous ops.
                    copy_stack_trace(node.outputs + switch_node.outputs, fct)
                    return fct
            except NotScalarConstantError:
                pass
            try:
                if (
                    get_scalar_constant_value(
                        switch_node.inputs[2], only_process_constants=True
                    )
                    == 0.0
                ):
                    listmul = node.inputs[:idx] + node.inputs[idx + 1 :]
                    fmul = mul(*(listmul + [switch_node.inputs[1]]))
                    # Copy over stacktrace for elementwise multiplication op
                    # from previous elementwise multiplication op.
                    # An error in the multiplication (e.g. errors due to
                    # inconsistent shapes), will point to the
                    # multiplication op.
                    copy_stack_trace(node.outputs, fmul)

                    fct = [switch(switch_node.inputs[0], fmul, 0)]
                    fct[0].tag.values_eq_approx = values_eq_approx_remove_nan

                    # Copy over stacktrace for switch op from both previous
                    # elementwise multiplication op and previous switch op,
                    # because an error in this part can be caused by either
                    # of the two previous ops.
                    copy_stack_trace(node.outputs + switch_node.outputs, fct)
                    return fct
            except NotScalarConstantError:
                pass
    return False


@register_canonicalize
@local_optimizer([true_div, int_div])
def local_div_switch_sink(fgraph, node):
    """
    This optimization makes the following changes in the graph:

    ``at.div(at.switch(cond, 0, iff), A)`` -> ``at.switch(cond, 0, at.div(iff, A))``
    ``at.div(at.switch(cond, ift, 0), A)`` -> ``at.switch(cond, at.div(ift, A), 0)``

    where ``A`` is a symbolic variable.

    This is useful because ``A`` may not be numerically stable and give
    ``nan`` or ``inf`` values for cases where the switch returns 0.

    See `local_mul_switch_sink` for more details.

    """
    if node.op != true_div and node.op != int_div:
        return False
    op = node.op
    if node.inputs[0].owner and node.inputs[0].owner.op == switch:
        switch_node = node.inputs[0].owner
        try:
            if (
                get_scalar_constant_value(
                    switch_node.inputs[1], only_process_constants=True
                )
                == 0.0
            ):
                fdiv = op(switch_node.inputs[2], node.inputs[1])
                # Copy over stacktrace for elementwise division op
                # from previous elementwise multiplication op.
                # An error in the division (e.g. errors due to
                # inconsistent shapes or division by zero),
                # will point to the new division op.
                copy_stack_trace(node.outputs, fdiv)

                fct = [switch(switch_node.inputs[0], 0, fdiv)]
                fct[0].tag.values_eq_approx = values_eq_approx_remove_nan

                # Copy over stacktrace for switch op from both previous
                # elementwise division op and previous switch op,
                # because an error in this part can be caused by either
                # of the two previous ops.
                copy_stack_trace(node.outputs + switch_node.outputs, fct)
                return fct
        except NotScalarConstantError:
            pass
        try:
            if (
                get_scalar_constant_value(
                    switch_node.inputs[2], only_process_constants=True
                )
                == 0.0
            ):
                fdiv = op(switch_node.inputs[1], node.inputs[1])
                # Copy over stacktrace for elementwise division op
                # from previous elementwise multiplication op.
                # An error in the division (e.g. errors due to
                # inconsistent shapes or division by zero),
                # will point to the new division op.
                copy_stack_trace(node.outputs, fdiv)

                fct = [switch(switch_node.inputs[0], fdiv, 0)]
                fct[0].tag.values_eq_approx = values_eq_approx_remove_nan

                # Copy over stacktrace for switch op from both previous
                # elementwise division op and previous switch op,
                # because an error in this part can be caused by either
                # of the two previous ops.
                copy_stack_trace(node.outputs + switch_node.outputs, fct)
                return fct
        except NotScalarConstantError:
            pass
    return False


class AlgebraicCanonizer(LocalOptimizer):
    r"""Simplification tool.

    The variable is a ``local_optimizer``. It is best used
    with a ``TopoOptimizer`` in ``in_to_out`` order.

    Usage: ``AlgebraicCanonizer(main, inverse, reciprocal, calculate)``

    Parameters
    ----------
    main
        A suitable ``Op`` class that is commutative, associative and
        takes one to an arbitrary number of inputs, e.g. add or
        mul
    inverse
        An ``Op`` class such that ``inverse(main(x, y), y) == x``
        e.g. ``sub`` or true_div
    reciprocal
        A function such that ``main(x, reciprocal(y)) == inverse(x, y)``
        e.g. ``neg`` or ``reciprocal``
    calculate
        Function that takes a list of numpy.ndarray instances
        for the numerator, another list for the denumerator,
        and calculates ``inverse(main(\*num), main(\*denum))``. It
        takes a keyword argument, aslist. If True, the value
        should be returned as a list of one element, unless
        the value is such that value = main(). In that case,
        the return value should be an empty list.

    Examples
    --------
    >>> import aesara.tensor as at
    >>> from aesara.tensor.math_opt import AlgebraicCanonizer
    >>> add_canonizer = AlgebraicCanonizer(add, sub, neg, \\
    ...                                    lambda n, d: sum(n) - sum(d))
    >>> mul_canonizer = AlgebraicCanonizer(mul, true_div, inv, \\
    ...                                    lambda n, d: prod(n) / prod(d))

    Examples of optimizations ``mul_canonizer`` can perform:

    | x / x -> 1
    | (x * y) / x -> y
    | x / y / x -> 1 / y
    | x / y / z -> x / (y * z)
    | x / (y / z) -> (x * z) / y
    | (a / b) * (b / c) * (c / d) -> a / d
    | (2.0 * x) / (4.0 * y) -> (0.5 * x) / y
    | 2 * x / 2 -> x
    | x * y * z -> Elemwise(mul){x,y,z} #only one pass over the memory.
    |           !-> Elemwise(mul){x,Elemwise(mul){y,z}}

    """

    def __init__(self, main, inverse_fn, reciprocal_fn, calculate, use_reciprocal=True):
        self.main = main
        self.inverse = inverse_fn
        self.reciprocal = reciprocal_fn
        self.calculate = calculate
        self.use_reciprocal = use_reciprocal

        self.external_simplifiers = []

    def add_simplifier(self, simplifier, reason):
        self.external_simplifiers.append((reason, simplifier))

    def tracks(self):
        return [self.main, self.inverse, self.reciprocal]

    def get_num_denum(self, inp):
        r"""
        This extract two lists, ``num`` and ``denum``, such that the input is:
        ``self.inverse(self.main(\*num), self.main(\*denum))``. It returns
        the two lists in a ``(num, denum)`` pair.

        For example, for main, inverse and ``reciprocal = \*, / and inv()``,

        | input -> returned value (num, denum)

        | x*y -> ([x, y], [])
        | inv(x) -> ([], [x])
        | inv(x) * inv(y) -> ([], [x, y])
        | x*y/z -> ([x, y], [z])
        | log(x) / y * (z + x) / y -> ([log(x), z + x], [y, y])
        | (((a / b) * c) / d) -> ([a, c], [b, d])
        | a / (b / c) -> ([a, c], [b])
        | log(x) -> ([log(x)], [])
        | x**y -> ([x**y], [])
        | x * y * z -> ([x, y, z], [])

        """
        # This function is recursive.  The idea is that there is a
        # get_num_denum recursion in which the internal ops are all
        # one of (main, inverse, reciprocal, DimShuffle) and the
        # internal data nodes all have the dtype of the 'input'
        # argument. The leaf-Variables of the graph covered by the
        # recursion may be of any Variable type.

        if inp.owner is None or inp.owner.op not in [
            self.main,
            self.inverse,
            self.reciprocal,
        ]:
            if inp.owner and isinstance(inp.owner.op, DimShuffle):
                # If input is a DimShuffle of some input which does
                # something like this:

                # * change a vector of length N into a 1xN row matrix
                # * change a scalar into a 1x1x1 tensor
                # * in general, complete the shape of a tensor
                #   with broadcastable 1s to the *left*
                # Then we will simply discard the DimShuffle and return
                # the num/denum of its input
                dsn = inp.owner  # dimshuffle node
                dsop = dsn.op  # dimshuffle op

                # the first input of the dimshuffle i.e. the ndarray to redim
                dsi0 = dsn.inputs[0]

                # The compatible order is a DimShuffle "new_order" of the form:
                # ('x', ..., 'x', 0, 1, 2, ..., dimshuffle_input.type.ndim)

                # That kind of DimShuffle only adds broadcastable
                # dimensions on the left, without discarding any
                # existing broadcastable dimension and is inserted
                # automatically by Elemwise when the inputs have
                # different numbers of dimensions (hence why we can
                # discard its information - we know we can retrieve it
                # later on).
                compatible_order = ("x",) * (inp.type.ndim - dsi0.type.ndim) + tuple(
                    range(dsi0.type.ndim)
                )
                if dsop.new_order == compatible_order:
                    # If the "new_order" is the one we recognize,
                    # we return the num_denum of the dimshuffled input.
                    return self.get_num_denum(inp.owner.inputs[0])
                else:
                    # This is when the input isn't produced by main,
                    # inverse or reciprocal.
                    return [inp], []
            else:
                return [inp], []
        num = []
        denum = []
        parent = inp.owner

        # We get the (num, denum) pairs for each input
        # pairs = [self.get_num_denum(input2) if input2.type.dtype ==
        # input.type.dtype else ([input2], []) for input2 in
        # parent.inputs]
        pairs = [self.get_num_denum(input2) for input2 in parent.inputs]

        if parent.op == self.main:
            # If we have main(x, y, ...), numx, denumx, numy, denumy, ...
            # then num is concat(numx, numy, num...) and denum is
            # concat(denumx, denumy, denum...) note that main() can have any
            # number of arguments >= 0 concat is list concatenation
            num = reduce(list.__iadd__, map(operator.itemgetter(0), pairs))
            denum = reduce(list.__iadd__, map(operator.itemgetter(1), pairs))
        elif parent.op == self.inverse:
            # If we have inverse(x, y), numx, denumx, numy and denumy
            # then num is concat(numx, denumy) and denum is
            # concat(denumx, numy) note that inverse() is binary
            num = pairs[0][0] + pairs[1][1]
            denum = pairs[0][1] + pairs[1][0]
        elif parent.op == self.reciprocal:
            # If we have reciprocal(x), numx, denumx
            # then num is denumx and denum is numx
            # note that reciprocal() is unary
            num = pairs[0][1]
            denum = pairs[0][0]
        return num, denum

    def merge_num_denum(self, num, denum):
        r"""
        Utility function which takes two lists, num and denum, and
        returns something which is equivalent to inverse(main(\*num),
        main(\*denum)), but depends on the length of num and the length
        of denum (in order to minimize the number of operations).

        Let n = len(num) and d = len(denum):

        | n=0, d=0: neutral element (given by self.calculate([], []))
        |           (for example, this would be 0 if main is addition
        |           and 1 if main is multiplication)
        | n=1, d=0: num[0]
        | n=0, d=1: reciprocal(denum[0])
        | n=1, d=1: inverse(num[0], denum[0])
        | n=0, d>1: reciprocal(main(\*denum))
        | n>1, d=0: main(\*num)
        | n=1, d>1: inverse(num[0], main(\*denum))
        | n>1, d=1: inverse(main(\*num), denum[0])
        | n>1, d>1: inverse(main(\*num), main(\*denum))

        Given the values of n and d to which they are associated, all
        of the above are equivalent to:
        inverse(main(\*num), main(\*denum))

        """

        ln, ld = len(num), len(denum)
        if not ln and not ld:
            return as_tensor_variable(self.calculate([], []))
        if not ln:
            if self.use_reciprocal:
                return self.reciprocal(self.merge_num_denum(denum, []))
            else:
                ln = [self.calculate([], [], aslist=False)]
        if not ld:
            if ln == 1:
                # num[0] should always be a variable
                assert isinstance(num[0], Variable)
                return num[0]
            else:
                return self.main(*num)
        return self.inverse(
            self.merge_num_denum(num, []), self.merge_num_denum(denum, [])
        )

    def simplify(self, num, denum, out_type):
        """
        Shorthand for:

        .. code-block:: python

            self.simplify_constants(*self.simplify_factors(num, denum))

        """
        rval = self.simplify_constants(
            *self.simplify_factors(num, denum), out_type=out_type
        )
        for reason, simplifier in self.external_simplifiers:
            # TODO: document that 'reason' is associated with this
            #       simplification to help auditing when things go
            #       wrong
            rval = simplifier(*rval)
        return rval

    def simplify_factors(self, num, denum):
        """
        For any Variable r which is both in num and denum, removes it
        from both lists. Modifies the lists inplace. Returns the
        modified lists. For example:

        | [x], [x] -> [], []
        | [x, y], [x] -> [y], []
        | [a, b], [c, d] -> [a, b], [c, d]

        """
        ln = len(num)
        ld = len(denum)
        if ld > 2 and ln > 2:
            # Faster version for "big" inputs.
            while True:
                s = set(num)
                # Inputs can appear multiple times
                redo = len(s) != len(num)
                inter = s.intersection(denum)
                for v in inter:
                    num.remove(v)
                    denum.remove(v)
                if not redo or not inter:
                    break
        else:
            for v in list(num):
                if v in denum:
                    num.remove(v)
                    denum.remove(v)
        return num, denum

    def simplify_constants(self, orig_num, orig_denum, out_type=None):
        """
        Find all constants and put them together into a single constant.

        Finds all constants in orig_num and orig_denum (using
        get_constant) and puts them together into a single
        constant. The constant is inserted as the first element of the
        numerator. If the constant is the neutral element, it is
        removed from the numerator.

        Examples
        --------
        Let main be multiplication:

        | [2, 3, x], [] -> [6, x], []
        | [x, y, 2], [4, z] -> [0.5, x, y], [z]
        | [x, 2, y], [z, 2] -> [x, y], [z]

        """
        # Lists representing the numerator and denumerator
        num, denum = [], []

        # Lists representing the *constant* elements of num and denum
        numct, denumct = [], []

        for v in orig_num:
            ct = get_constant(v)
            if ct is not None:
                # We found a constant in the numerator!
                # We add it to numct
                numct.append(ct)
            else:
                num.append(v)
        for v in orig_denum:
            ct = get_constant(v)
            if ct is not None:
                denumct.append(ct)
            else:
                denum.append(v)

        if self.use_reciprocal or num:
            # This will calculate either:
            # [inverse(main(*numct), main(*denumct))]
            # [] - if inverse(main(*numct), main(*denumct)) is the
            # neutral element
            ct = self.calculate(numct, denumct, aslist=True, out_type=out_type)
        else:
            # This happens if we don't allow the reciprocal and the
            # numerator is empty. That means we will need to represent
            # reciprocal(x) like inverse(neutral_element, x) so
            # we can't allow ct == []
            # TODO: why is this branch needed when merge_num_denum
            # does it for us?
            ct = [self.calculate(numct, denumct, aslist=False, out_type=out_type)]

        # Wrapping ct in a Constant with the right dtype
        ct = [constant(c, dtype=out_type.dtype) for c in ct]

        if orig_num and len(numct) == 1 and len(denumct) == 0 and ct:
            # In that case we should only have one constant in `ct`.
            assert len(ct) == 1
            first_num_ct = get_constant(orig_num[0])
            if first_num_ct is not None and ct[0].type.values_eq(
                ct[0].data, first_num_ct
            ):
                # This is an important trick :( if it so happens that:
                # * there's exactly one constant on the numerator and none on
                #   the denominator
                # * it's not the neutral element (ct is an empty list in that
                #   case)
                # * the constant is the same as the first argument in the
                #   numerator (we only check the first argument because the
                #   canonizer puts the computed constants first)
                # -> then we return very exactly the original num/denum.
                # If we don't do that the optimizer will just loop
                # infinitely because it will not catch on that there are
                # no changes to be made and every time it will want to
                # replace something by the same thing...
                # Note that it is important to use `values_eq` instead of
                # the == operator, to handle NaN values correctly.
                return orig_num, orig_denum

        return ct + num, denum

    def transform(self, fgraph, node):
        op = node.op
        if op not in [self.main, self.inverse, self.reciprocal]:
            return False

        assert len(node.outputs) == 1
        out = node.outputs[0]

        out_clients = fgraph.clients.get(out)

        if not out_clients:
            return False

        # check if any of the clients of this node would be part of
        # this canonized graph...  if so, we do nothing and wait for
        # them to be transformed.
        for c, c_idx in out_clients:
            if c == "output":
                continue
            while (
                isinstance(getattr(c, "op", None), DimShuffle)
                and len(fgraph.clients[c.outputs[0]]) <= 1
            ):
                c = fgraph.clients[c.outputs[0]][0][0]
            if getattr(c, "op", "") in [self.main, self.inverse, self.reciprocal]:
                return False

        # Here we make the canonical version of the graph around this node
        # See the documentation of get_num_denum and simplify
        orig_num, orig_denum = self.get_num_denum(node.outputs[0])
        num, denum = self.simplify(list(orig_num), list(orig_denum), out.type)

        def same(x, y):
            return len(x) == len(y) and all(np.all(xe == ye) for xe, ye in zip(x, y))

        if (
            same(orig_num, num)
            and same(orig_denum, denum)
            and
            # Check to see if we've collapsed some nested ops.
            not (
                len(orig_denum) == 0
                and
                # Make sure this change would increase the number of vector
                # arguments--decreasing the number of unnecessary `self.main`
                # nodes.
                len(node.inputs) < len(orig_num)
            )
            and
            # Do a similar check for the reciprocal op.
            not (
                self.use_reciprocal
                and node.op == self.reciprocal
                and len(orig_num) == 0
                and node.inputs[0].owner
                and len(node.inputs[0].owner.inputs) < len(orig_denum)
            )
        ):
            return False

        new = self.merge_num_denum(num, denum)
        if new.type.dtype != out.type.dtype:
            new = cast(new, out.type.dtype)

        if new.type != out.type:
            new = fill_chain(new, node.inputs)[0]

        if new.type == out.type:
            new.tag.values_eq_approx = values_eq_approx_remove_inf_nan
            copy_stack_trace(out, new)
            return [new]
        else:
            return False

    def __str__(self):
        return getattr(
            self,
            "name",
            f"AlgebraicCanonizer({self.main}, {self.inverse}, {self.reciprocal})",
        )


def mul_calculate(num, denum, aslist=False, out_type=None):
    if not num and not denum:
        # Smallest 1 possible.
        if aslist:
            return []
        else:
            return np.int8(1)

    # Make sure we do not accidentally upcast data types.
    if out_type is None:
        out_dtype = aes.upcast(*[v.dtype for v in (num + denum)])
    else:
        out_dtype = out_type.dtype
    one = _asarray(1, dtype=out_dtype)

    v = reduce(np.multiply, num, one) / reduce(np.multiply, denum, one)
    if aslist:
        if np.all(v == 1):
            return []
        else:
            return [v]
    return v


local_mul_canonizer = AlgebraicCanonizer(
    mul, true_div, reciprocal, mul_calculate, False
)
register_canonicalize(local_mul_canonizer, name="local_mul_canonizer")


@register_canonicalize
@local_optimizer([neg])
def local_neg_to_mul(fgraph, node):
    if node.op == neg:
        return [mul(np.array(-1, dtype=node.inputs[0].dtype), node.inputs[0])]


@register_specialize
@local_optimizer([Sum, Prod])
def local_sum_prod_mul_by_scalar(fgraph, node):
    """
    sum(scalar * smth) -> scalar * sum(smth)
    sum(-smth) -> -sum(smth)

    or

    prod(scalar * smth) -> scalar ** size(smth) * prod(smth)
    prod(-smth) -> -1 ** size(smth) * prod(smth)

    """
    # TODO: if the the thing inside the Sum is a division,
    # we should get at the numerator....
    if isinstance(node.op, (Sum, Prod)):
        (node_inps,) = node.inputs
        if node_inps.owner and node_inps.owner.op == mul:
            terms = node_inps.owner.inputs
            scalars = [t.dimshuffle() for t in terms if np.all(t.type.broadcastable)]

            if len(scalars) == 0:
                # Nothing to optimize here
                return

            non_scalars = [t for t in terms if not np.all(t.broadcastable)]

            # Perform the op only on the non-scalar inputs, if applicable
            if len(non_scalars) == 0:
                new_op_input_nb_elements = 1
                new_op_output = 1
            elif len(non_scalars) == 1:
                new_op_input_nb_elements = non_scalars[0].size
                new_op_output = node.op(non_scalars[0])
            else:
                new_op_input = mul(*non_scalars)
                # We assume that errors always come from the prod/mul op in the
                # original computational graph, and therefore need to only
                # copy over its output stacktrace.
                copy_stack_trace(node.outputs, new_op_input)

                new_op_input_nb_elements = new_op_input.size
                new_op_output = node.op(new_op_input)

            if len(non_scalars) != 0:
                # Copy over stacktrace from previous output to new mul op,
                # for same reason as above.
                copy_stack_trace(node.outputs, new_op_output)

            # If `node.op` is a `Prod`, then the scalars need to be raised to
            # the power of the number of elements in the input to the `Prod`
            if isinstance(node.op, Prod) and new_op_input_nb_elements != 1:

                scalars = [s ** new_op_input_nb_elements for s in scalars]

            # Scale the output of the op by the scalars and return as
            # replacement for the original output
            mul_inputs = scalars
            if new_op_input_nb_elements != 1:
                mul_inputs.append(new_op_output)

            if len(mul_inputs) == 1:
                # Copy over stacktrace from previous output to new mul op,
                # for same reason as above.
                copy_stack_trace(node.outputs, mul_inputs)

                return mul_inputs
            else:
                ret = mul(*mul_inputs)
                # Copy over stacktrace from previous output to new mul op,
                # for same reason as above.
                copy_stack_trace(node.outputs, [ret] + mul_inputs)

                return [ret]

        if isinstance(node.op, Sum) and node_inps.owner and node_inps.owner.op == neg:
            s = node.op(node_inps.owner.inputs[0])
            ret = neg(s)
            # There are never errors in the negative op, thus
            # we need only to copy over stacktrace from previous output node to
            # the two new ops.
            copy_stack_trace(node.outputs, [s, ret])

            return [ret]


@register_specialize
@local_optimizer([Elemwise])
def local_elemwise_sub_zeros(fgraph, node):
    """
    Elemwise{sub}(X,X) -> zeros_like(X)
    """
    if (
        isinstance(node.op, Elemwise)
        and node.op.scalar_op.nin == 2
        and node.op.scalar_op == aes.sub
        and node.inputs[0] == node.inputs[1]
    ):
        res = zeros_like(node.inputs[0])
        # Copy over stacktrace from previous output.
        # This could help for failures due to out-of-memory.
        copy_stack_trace(node.outputs, res)
        return [res]


@register_useless
@register_specialize
@register_stabilize
@register_canonicalize
@local_optimizer([Elemwise])
def local_useless_elemwise_comparison(fgraph, node):
    """...

    :note: These cases appear in the graph generated by scan.
           These optimizations will make the graph easier to read.
    # Comparing to itself is constant
    Elemwise[{LT,GT}](X, X) -> Elemwise[zeros](X)
    Elemwise[{LE,GE}](X, X) -> Elemwise[ones](X)
    Elemwise[{minimum,maximum}](X, X) -> X

    # Comparing shape to 0 can be constant
    Elemwise[LT](X.shape[i], 0) -> Elemwise[zeros](X)
    Elemwise[GE](X.shape[i], 0) -> Elemwise[ones](X)
    Elemwise[maximum](X.shape[i], 0) -> X.shape[i]
    Elemwise[maximum](0, X.shape[i]) -> X.shape[i]
    Elemwise[minimum](X.shape[i], 0) -> 0
    Elemwise[minimum](0, X.shape[i]) -> 0

    # The shape can be replaced with sum of shapes
    Elemwise[LT](add([anything that is shapes]), 0) -> Elemwise[zeros](X)
    Elemwise[GE](add([anything that is shapes]), 0) -> Elemwise[ones](X)

    # Shapes are never negative
    # Needed by Reshape.infer_shape
    Elemwise[EQ](Subtensor(Shape(x)), -N) -> Elemwise[zeros](X)

    """
    if not isinstance(node.op, Elemwise):
        return
    if node.op.scalar_op.nin != 2:
        return

    # We call zeros_like and one_like with opt=True to generate a
    # cleaner graph.
    dtype = node.outputs[0].dtype

    # Elemwise[{LT,GT}](X, X) -> Elemwise[zeros](X)
    if (
        isinstance(node.op.scalar_op, (aes.LT, aes.GT))
        and node.inputs[0] is node.inputs[1]
    ):
        res = zeros_like(node.inputs[0], dtype=dtype, opt=True)
        # Copy over stacktrace from previous output.
        copy_stack_trace(node.outputs, res)
        return [res]
    # Elemwise[{LE,GE}](X, X) -> Elemwise[ones](X)
    if (
        isinstance(node.op.scalar_op, (aes.LE, aes.GE))
        and node.inputs[0] is node.inputs[1]
    ):
        res = ones_like(node.inputs[0], dtype=dtype, opt=True)

        # Copy over stacktrace from previous output.
        copy_stack_trace(node.outputs, res)
        return [res]
    # Elemwise[{minimum,maximum}](X, X) -> X
    if (
        isinstance(node.op.scalar_op, (aes.ScalarMinimum, aes.ScalarMaximum))
        and node.inputs[0] is node.inputs[1]
    ):
        res = node.inputs[0]
        # Copy over stacktrace from previous output.
        copy_stack_trace(node.outputs, res)
        return [res]

    # Elemwise[LT](X.shape[i], 0) -> Elemwise[zeros](X)
    if (
        isinstance(node.op.scalar_op, aes.LT)
        and node.inputs[0].owner
        and isinstance(node.inputs[0].owner.op, Shape_i)
        and extract_constant(node.inputs[1], only_process_constants=True) == 0
    ):
        res = zeros_like(node.inputs[0], dtype=dtype, opt=True)
        # Copy over stacktrace from previous output.
        copy_stack_trace(node.outputs, res)
        return [res]
    # Elemwise[GE](X.shape[i], 0) -> Elemwise[ones](X)
    if (
        isinstance(node.op.scalar_op, aes.GE)
        and node.inputs[0].owner
        and isinstance(node.inputs[0].owner.op, Shape_i)
        and extract_constant(node.inputs[1], only_process_constants=True) == 0
    ):
        res = ones_like(node.inputs[0], dtype=dtype, opt=True)
        # Copy over stacktrace from previous output.
        copy_stack_trace(node.outputs, res)
        return [res]
    # Elemwise[maximum](X.shape[i], 0) -> X.shape[i]
    if (
        isinstance(node.op.scalar_op, aes.ScalarMaximum)
        and node.inputs[0].owner
        and isinstance(node.inputs[0].owner.op, Shape_i)
        and extract_constant(node.inputs[1], only_process_constants=True) == 0
    ):
        # No need to copy over stacktrace.
        return [node.inputs[0]]
    # Elemwise[maximum](0, X.shape[i]) -> X.shape[i]
    if (
        isinstance(node.op.scalar_op, aes.ScalarMaximum)
        and extract_constant(node.inputs[0], only_process_constants=True) == 0
        and node.inputs[1].owner
        and isinstance(node.inputs[1].owner.op, Shape_i)
    ):
        # No need to copy over stacktrace.
        return [node.inputs[1]]
    # Elemwise[minimum](X.shape[i], 0) -> 0
    if (
        isinstance(node.op.scalar_op, aes.ScalarMinimum)
        and node.inputs[0].owner
        and isinstance(node.inputs[0].owner.op, Shape_i)
        and extract_constant(node.inputs[1], only_process_constants=True) == 0
    ):
        res = zeros_like(node.inputs[0], dtype=dtype, opt=True)
        # Copy over stacktrace from previous output.
        copy_stack_trace(node.outputs, res)
        return [res]

    # Elemwise[minimum](0, X.shape[i]) -> 0
    if (
        isinstance(node.op.scalar_op, aes.ScalarMinimum)
        and extract_constant(node.inputs[0], only_process_constants=True) == 0
        and node.inputs[1].owner
        and isinstance(node.inputs[1].owner.op, Shape_i)
    ):
        res = zeros_like(node.inputs[1], dtype=dtype, opt=True)
        # Copy over stacktrace from previous output.
        copy_stack_trace(node.outputs, res)
        return [res]

    # Elemwise[LT](add([anything that is shapes]), 0) -> Elemwise[zeros](X)
    if (
        isinstance(node.op.scalar_op, aes.LT)
        and node.inputs[0].owner
        and isinstance(node.inputs[0].owner.op, Elemwise)
        and isinstance(node.inputs[0].owner.op.scalar_op, aes.Add)
        and all(
            [
                isinstance(var.owner and var.owner.op, Shape_i)
                for var in node.inputs[0].owner.inputs
            ]
        )
        and extract_constant(node.inputs[1], only_process_constants=True) == 0
    ):
        res = zeros_like(node.inputs[0], dtype=dtype, opt=True)
        # Copy over stacktrace from previous output.
        copy_stack_trace(node.outputs, res)
        return [res]
    # Elemwise[GE](add([anything that is shapes]), 0) -> Elemwise[ones](X)
    if (
        isinstance(node.op.scalar_op, aes.GE)
        and node.inputs[0].owner
        and isinstance(node.inputs[0].owner.op, Elemwise)
        and isinstance(node.inputs[0].owner.op.scalar_op, aes.Add)
        and all(
            [
                isinstance(var.owner and var.owner.op, Shape_i)
                for var in node.inputs[0].owner.inputs
            ]
        )
        and extract_constant(node.inputs[1], only_process_constants=True) == 0
    ):
        res = ones_like(node.inputs[0], dtype=dtype, opt=True)

        # Copy over stacktrace from previous output.
        copy_stack_trace(node.outputs, res)
        return [res]

        # Elemwise[EQ](Subtensor(Shape(x)), -N)
        # Elemwise[EQ](somegraph that only depend of shape, -N)
        # TODO: handle the case where the -N is on either side
        """
 |Elemwise{eq,no_inplace} [id B] ''
 | |Subtensor{int64} [id C] ''
 | | |Join [id D] ''
 | | | |TensorConstant{0} [id E]
 | | | |Subtensor{int64:int64:} [id F] ''
 | | | | |Shape [id G] ''
        """

    def investigate(node):
        " Return True if values will be shapes, so >= 0"
        if isinstance(node.op, (Shape, Shape_i)):
            return True
        elif isinstance(node.op, Subtensor) and node.inputs[0].owner:
            return investigate(node.inputs[0].owner)
        elif isinstance(node.op, Join):
            return all(v.owner and investigate(v.owner) for v in node.inputs[1:])
        elif isinstance(node.op, MakeVector):
            return all(v.owner and investigate(v.owner) for v in node.inputs)

    if (
        isinstance(node.op.scalar_op, aes.EQ)
        and node.inputs[0].owner
        and investigate(node.inputs[0].owner)
    ):
        try:
            cst = get_scalar_constant_value(node.inputs[1], only_process_constants=True)

            res = zeros_like(node.inputs[0], dtype=dtype, opt=True)

            if cst < 0:
                # Copy over stacktrace from previous output.
                copy_stack_trace(node.outputs, res)

                return [res]

        except NotScalarConstantError:
            pass
    return


@register_canonicalize
@register_specialize
@local_optimizer([Sum, Prod])
def local_sum_prod_div_dimshuffle(fgraph, node):
    """
    sum(a / dimshuffle{...}(b), axis=l) -> sum(a, axis={...}) / b,
    if dimension l of the DimShuffle is 'x'

    or

    prod(a / dimshuffle{...}(b), axis=l) ->
    prod(a, axis={...}) / b ** a.shape[l],
    if dimension l of the DimShuffle is 'x'
    """

    # It does not make much sense now to extend it to the case where the
    # dimshuffle is in the numerator, since elemwise inversion of the
    # denominator would still be needed before the summation or production.

    if isinstance(node.op, (Sum, Prod)):
        axis = node.op.axis
        if axis is None:
            axis = list(range(node.inputs[0].ndim))
        node_input = node.inputs[0]
        if node_input.owner and node_input.owner.op == true_div:
            numerator, denominator = node_input.owner.inputs

            if denominator.owner and isinstance(denominator.owner.op, DimShuffle):
                dimshuffle_input = denominator.owner.inputs[0]
                dimshuffle_order = denominator.owner.op.new_order

                compatible_dims = []
                incompatible_dims = []
                for ax in axis:
                    if ax < len(dimshuffle_order) and dimshuffle_order[ax] == "x":
                        compatible_dims.append(ax)
                    else:
                        incompatible_dims.append(ax)
                reordered_incompatible_dims = []
                for ic_ax in incompatible_dims:
                    reordered_incompatible_dims.append(
                        ic_ax - sum([1 for c_ax in compatible_dims if c_ax < ic_ax])
                    )

                if len(compatible_dims) > 0:
                    optimized_dimshuffle_order = list(
                        ax
                        for i, ax in enumerate(dimshuffle_order)
                        if (i not in axis) or (ax != "x")
                    )

                    # Removing leading 'x' (since it will be done automatically)
                    while (
                        len(optimized_dimshuffle_order) > 0
                        and optimized_dimshuffle_order[0] == "x"
                    ):
                        del optimized_dimshuffle_order[0]

                    # if optimized_dimshuffle_order is sorted with
                    # not 'x', then dimshuffle is useless.
                    if all(i == e for i, e in enumerate(optimized_dimshuffle_order)):
                        optimized_dimshuffle = dimshuffle_input
                    else:
                        optimized_dimshuffle = DimShuffle(
                            dimshuffle_input.type.broadcastable,
                            optimized_dimshuffle_order,
                        )(dimshuffle_input)

                    if isinstance(node.op, Sum):
                        op_on_compatible_dims = at_sum(numerator, axis=compatible_dims)
                        rval = true_div(op_on_compatible_dims, optimized_dimshuffle)
                        if len(reordered_incompatible_dims) > 0:
                            rval = at_sum(rval, axis=reordered_incompatible_dims)
                    elif isinstance(node.op, Prod):
                        op_on_compatible_dims = prod(numerator, axis=compatible_dims)
                        dtype = numerator.dtype
                        rval = true_div(
                            op_on_compatible_dims,
                            (
                                optimized_dimshuffle
                                ** prod(
                                    [
                                        numerator.shape[ax].astype(dtype)
                                        for ax in compatible_dims
                                    ]
                                )
                            ),
                        )
                        if len(reordered_incompatible_dims) > 0:
                            rval = prod(rval, axis=reordered_incompatible_dims)
                    return [rval]


@register_canonicalize
@local_optimizer([Sum, Prod])
def local_sum_prod_all_to_none(fgraph, node):
    """
    Sum{0,1,...N} -> Sum{} or
    Prod{0,1,...N} -> Prod{}

    """
    if isinstance(node.op, Sum) or isinstance(node.op, Prod):
        opt_type = Sum if isinstance(node.op, Sum) else Prod
        # if all the axes are named, then use None as a shorthand
        # this permits more merging
        if node.op.axis is None:
            return
        if set(node.op.axis) == set(range(node.inputs[0].type.ndim)):
            return [opt_type(axis=None, dtype=node.op.dtype)(node.inputs[0])]


@register_canonicalize
@local_optimizer([Sum, Prod])
def local_op_of_op(fgraph, node):
    """
    Prod(Prod()) -> single Prod()
    or
    Sum(Sum()) -> single Sum()

    """
    if isinstance(node.op, Prod) or isinstance(node.op, Sum):
        opt_type = Sum if isinstance(node.op, Sum) else Prod
        (node_inps,) = node.inputs
        out_dtype = node.op.dtype
        # We manipulate the graph so this is done to make sure the opt
        # doesn't affect other computations.
        if len(fgraph.clients[node_inps]) == 1:
            if node_inps.owner and (isinstance(node_inps.owner.op, node.op.__class__)):

                # check to see either the inner or outer prod is doing a
                # product over all axis, in which case we can remove it
                if node_inps.owner.op.axis is None or node.op.axis is None:
                    return [opt_type(None, dtype=out_dtype)(node_inps.owner.inputs[0])]

                # figure out which axes were in the original sum
                newaxis = list(tuple(node_inps.owner.op.axis))
                for i in node.op.axis:
                    new_i = i
                    for ii in node_inps.owner.op.axis:
                        if new_i >= ii:
                            new_i += 1
                    assert new_i not in newaxis
                    newaxis.append(new_i)

                assert len(newaxis) == len(
                    list(node_inps.owner.op.axis) + list(node.op.axis)
                )

                combined = opt_type(newaxis, dtype=out_dtype)
                return [combined(node_inps.owner.inputs[0])]


ALL_REDUCE = (
    [
        CAReduce,
        All,
        Any,
        Sum,
        Prod,
        ProdWithoutZeros,
    ]
    + CAReduce.__subclasses__()
    + NonZeroCAReduce.__subclasses__()
)


@register_canonicalize
@register_uncanonicalize  # Needed for MaxAndArgmax -> CAReduce
@local_optimizer(ALL_REDUCE)
def local_reduce_join(fgraph, node):
    """
    Reduce{scalar.op}(Join(axis=0, a, b), axis=0) -> Elemwise{scalar.op}(a, b)

    Notes
    -----
    Supported scalar.op are Maximum, Minimum in some cases and Add and Mul in
    all cases.

    Currently we must reduce on axis 0. It is probably extensible to the case
    where we join and reduce on the same set of axis.

    """
    if (
        isinstance(node.op, CAReduce)
        and node.inputs[0].owner
        and isinstance(node.inputs[0].owner.op, Join)
    ):
        join_node = node.inputs[0].owner
        if extract_constant(join_node.inputs[0], only_process_constants=True) != 0:
            return

        if isinstance(node.op.scalar_op, (aes.ScalarMaximum, aes.ScalarMinimum)):
            # Support only 2 inputs for now
            if len(join_node.inputs) != 3:
                return
        elif not isinstance(node.op.scalar_op, (aes.Add, aes.Mul)):
            return
        elif len(join_node.inputs) <= 2:
            # This is a useless join, that will get removed by another opt.
            return

        new_inp = []
        for inp in join_node.inputs[1:]:
            inp = inp.owner
            if not inp:
                return
            if not isinstance(inp.op, DimShuffle) or inp.op.new_order != ("x",) + tuple(
                range(inp.inputs[0].ndim)
            ):
                return
            new_inp.append(inp.inputs[0])
        ret = Elemwise(node.op.scalar_op)(*new_inp)

        if ret.dtype != node.outputs[0].dtype:
            # The reduction do something about the dtype.
            return

        reduce_axis = node.op.axis
        if reduce_axis is None:
            reduce_axis = tuple(range(node.inputs[0].ndim))

        if len(reduce_axis) != 1 or 0 not in reduce_axis:
            return

        # We add the new check late to don't add extra warning.
        try:
            join_axis = get_scalar_constant_value(
                join_node.inputs[0], only_process_constants=True
            )

            if join_axis != reduce_axis[0]:
                return
        except NotScalarConstantError:
            return

        return [ret]


@register_canonicalize("fast_compile", "local_cut_useless_reduce")
@register_useless("local_cut_useless_reduce")
@local_optimizer(ALL_REDUCE)
def local_useless_reduce(fgraph, node):
    """Sum(a, axis=[]) -> a  """
    if isinstance(node.op, CAReduce):
        (summed,) = node.inputs
        # if reduce were doing anything, the output ndim would be reduced
        if summed.type == node.outputs[0].type:
            return [summed]


@register_canonicalize
@register_uncanonicalize
@register_specialize
@local_optimizer(ALL_REDUCE)
def local_reduce_broadcastable(fgraph, node):
    """Remove reduction over broadcastable dimensions."""
    if isinstance(node.op, CAReduce):
        (reduced,) = node.inputs
        odtype = node.outputs[0].dtype
        if node.op.axis is None:
            if all(reduced.broadcastable):
                return [reduced.dimshuffle().astype(odtype)]
        else:
            axis = list(node.op.axis)
            cuttable = [a for a in axis if reduced.broadcastable[a]]
            if cuttable:
                # -- we can remove some axes of summation,
                #    which simplifies the codegen for sum, especially on GPU
                new_axis = []
                pattern = []
                ii = 0
                for p in range(reduced.ndim):
                    if p not in cuttable:
                        if p in axis:
                            new_axis.append(ii)
                        pattern.append(p)
                        ii += 1
                new_reduced = reduced.dimshuffle(*pattern)
                if new_axis:
                    if type(node.op) == CAReduce:
                        # This case handles `CAReduce` instances
                        # (e.g. generated by `scalar_elemwise`), and not the
                        # scalar `Op`-specific subclasses
                        # TODO FIXME: This highlights a major design flaw in
                        # `CAReduce` (or at least our use of it), and it needs
                        # to be fixed
                        new_op = node.op.__class__(node.op.scalar_op, axis=new_axis)
                    else:
                        new_op = node.op.__class__(axis=new_axis)
                    return [new_op(new_reduced)]
                else:
                    # -- in this case we can remove the reduction completely
                    return [new_reduced.astype(odtype)]


@register_specialize
@local_optimizer([Sum, Prod])
def local_opt_alloc(fgraph, node):
    """
    sum(alloc(constant,shapes...)) => constant*prod(shapes)
    or
    prod(alloc(constant,shapes...)) => constant**prod(shapes)

    """
    if isinstance(node.op, Sum) or isinstance(node.op, Prod):
        (node_inps,) = node.inputs
        if node_inps.owner and isinstance(node_inps.owner.op, Alloc):
            inp = node_inps.owner.inputs[0]
            shapes = node_inps.owner.inputs[1:]
            try:
                val = get_scalar_constant_value(inp, only_process_constants=True)
                assert val.size == 1
                val = val.reshape(1)[0]
                # check which type of op
                size = mul(*shapes)
                if inp.dtype in ("float16", "float32"):
                    # shapes are ints and normally int64.
                    # We don't want to have a float64 upcast
                    # We don't want to downcast to float16
                    # as we fear it could loose too much precision
                    # that will be amplified by the mul/pow below.
                    size = size.astype("float32")
                if node.op.axis is None or node.op.axis == tuple(range(inp.ndim)):
                    if isinstance(node.op, Sum):
                        val = val * size
                    else:
                        val = val ** size
                    # Sum can change the input dtype (upcast or bool
                    # -> float32) by default or by user request.
                    # We can ignore the acc_dtype, as there is only 1
                    # elemwise we will do and not a sequence, so there is no
                    # accumulation of errors.
                    # So mostly, we just need to cast the output to the old
                    # dtype.
                    val = val.astype(node.outputs[0].dtype)
                    return [val]
                to_prod = [shapes[i] for i in range(len(shapes)) if i in node.op.axis]
                if to_prod:
                    size = mul(*to_prod)
                    if isinstance(node.op, Sum):
                        val *= size
                    else:
                        val = val ** size
                # See comments above.
                val = val.astype(node.outputs[0].dtype)
                return [
                    alloc(
                        val,
                        *[
                            shapes[i]
                            for i in range(len(shapes))
                            if i not in node.op.axis
                        ],
                    )
                ]
            except NotScalarConstantError:
                pass


@register_specialize
@local_optimizer([neg])
def local_neg_div_neg(fgraph, node):
    """
    - (-a / b) -> a / b

    Also performs - (c / b) -> ((-c) / b) when c is a scalar constant.

    """
    if node.op == neg:
        if node.inputs[0].owner and node.inputs[0].owner.op == true_div:
            frac = node.inputs[0]
            num, denom = frac.owner.inputs
            if num.owner and num.owner.op == neg:
                if len(fgraph.clients[frac]) == 1:
                    # No other clients of the original division
                    new_num = num.owner.inputs[0]
                    return [true_div(new_num, denom)]
            elif np.all(num.broadcastable) and isinstance(num, Constant):
                if len(fgraph.clients[frac]) == 1:
                    new_num = -num.data
                    return [true_div(new_num, denom)]


@register_canonicalize
@local_optimizer([mul])
def local_mul_zero(fgraph, node):
    """
    As part of canonicalization, we replace multiplication by zero
    with zero.

    """
    if node.op == mul:
        otype = node.outputs[0].type

        for i in node.inputs:
            try:
                value = get_scalar_constant_value(i)
            except NotScalarConstantError:
                continue
            # print 'MUL by value', value, node.inputs
            if value == 0:
                # print '... returning zeros'
                return fill_chain(_asarray(0, dtype=otype.dtype), node.inputs)


# TODO: Add this to the canonicalization to reduce redundancy.
@register_specialize
@local_optimizer([true_div])
def local_div_to_reciprocal(fgraph, node):
    if node.op == true_div and np.all(get_constant(node.inputs[0]) == 1.0):
        out = node.outputs[0]
        new_out = reciprocal(local_mul_canonizer.merge_num_denum(node.inputs[1:], []))
        # The ones could have forced upcasting
        if new_out.dtype != out.dtype:
            new_out = cast(new_out, dtype=out.dtype)
        # The ones could have forced a specific length
        if not out.type.is_super(new_out.type):
            new_out = broadcast_like(new_out, out, fgraph)
        return [new_out]
    else:
        return False


@register_canonicalize
@local_optimizer([reciprocal])
def local_reciprocal_canon(fgraph, node):
    if node.op == reciprocal:
        return [at_pow(node.inputs[0], -1.0)]
    else:
        return False


@register_canonicalize
@local_optimizer([at_pow])
def local_pow_canonicalize(fgraph, node):
    if node.op == at_pow:
        cst = get_constant(node.inputs[1])
        if cst == 0:
            return [broadcast_like(1, node.outputs[0], fgraph)]
        if cst == 1:
            return [broadcast_like(node.inputs[0], node.outputs[0], fgraph)]
    else:
        return False


@register_specialize
@local_optimizer([mul])
def local_mul_to_sqr(fgraph, node):
    """
    x*x -> sqr(x)

    This is faster on the GPU when memory fetching is a big part of
    the computation time.

    """
    if node.op == mul:
        if len(node.inputs) == 2:
            if node.inputs[0] is node.inputs[1]:
                return [sqr(node.inputs[0])]


@register_canonicalize
@local_optimizer([int_div])
def local_intdiv_by_one(fgraph, node):
    """x // 1 -> x"""
    if node.op in [int_div]:
        if isinstance(node.inputs[1], TensorConstant) and np.all(
            node.inputs[1].value == 1
        ):
            return [node.inputs[0].astype(node.outputs[0].dtype)]


@register_canonicalize
@register_specialize
@local_optimizer([int_div, true_div])
def local_zero_div(fgraph, node):
    """0 / x -> 0"""
    if isinstance(node.op, Elemwise) and isinstance(
        node.op.scalar_op, (aes.IntDiv, aes.TrueDiv)
    ):
        if get_constant(node.inputs[0]) == 0:
            ret = broadcast_like(0, node.outputs[0], fgraph)
            ret.tag.values_eq_approx = values_eq_approx_remove_nan
            return [ret]


@register_specialize
@local_optimizer([at_pow])
def local_pow_specialize(fgraph, node):
    # here, we are past the point of canonicalization, so we don't want
    # to put in un-necessary fills.
    if node.op == at_pow:
        # the idea here is that we have pow(x, y)
        odtype = node.outputs[0].dtype
        xsym = node.inputs[0]
        ysym = node.inputs[1]
        y = get_constant(ysym)
        if (y is not None) and encompasses_broadcastable(
            xsym.type.broadcastable, ysym.type.broadcastable
        ):
            rval = None

            if np.all(y == 2):
                rval = [sqr(xsym)]
            if np.all(y == 1):
                rval = [xsym]
            if np.all(y == 0):
                rval = [fill(xsym, np.asarray(1, dtype=odtype))]
            if np.all(y == 0.5):
                rval = [sqrt(xsym)]
            if np.all(y == -0.5):
                rval = [reciprocal(sqrt(xsym))]
            if np.all(y == -1):
                rval = [reciprocal(xsym)]
            if np.all(y == -2):
                rval = [reciprocal(sqr(xsym))]
            if rval:
                rval[0] = cast(rval[0], odtype)
                assert rval[0].type == node.outputs[0].type, (rval, node.outputs)
                return rval
    else:
        return False


@register_specialize_device
@local_optimizer([at_pow])
def local_pow_specialize_device(fgraph, node):
    """
    This optimization is not the same on all device. We do it only on cpu here.
    """
    if node.op == at_pow:
        # the idea here is that we have pow(x, y)
        odtype = node.outputs[0].dtype
        xsym = node.inputs[0]
        ysym = node.inputs[1]
        y = get_constant(ysym)

        # the next line is needed to fix a strange case that I don't
        # know how to make a separate test.
        # That happen in the test_opt.py:test_log_erfc test.
        # y is a ndarray with dtype int8 and value 2,4 or 6. This make
        # the abs(y) <= 512 fail!
        # taking the value outside ndarray solve the problem.
        # it could be that in that case, numpy make the comparison
        # into the wrong type(do in int8 that overflow.)
        if isinstance(y, np.ndarray):
            assert y.size == 1
            try:
                y = y[0]
            except IndexError:
                pass
        if (y is not None) and encompasses_broadcastable(
            xsym.type.broadcastable, ysym.type.broadcastable
        ):
            rval = None
            # 512 is too small for the cpu and too big for some gpu!
            if abs(y) == int(abs(y)) and abs(y) <= 512:
                pow2 = [xsym]
                pow2_scal = [aes.get_scalar_type(xsym.dtype)()]
                y_to_do = abs(y)
                for i in range(int(np.log2(y_to_do))):
                    pow2.append(sqr(pow2[i]))
                    pow2_scal.append(aes.sqr(pow2_scal[i]))
                rval1 = None
                rval1_scal = None
                while y_to_do > 0:
                    log_to_do = int(np.log2(y_to_do))
                    if rval1:
                        rval1 *= pow2[log_to_do]
                        rval1_scal *= pow2_scal[log_to_do]
                    else:
                        rval1 = pow2[log_to_do]
                        rval1_scal = pow2_scal[log_to_do]
                    y_to_do -= 2 ** log_to_do

                if abs(y) > 2:
                    # We fuse all the pow together here to make
                    # compilation faster
                    rval1 = Elemwise(
                        aes.Composite([pow2_scal[0]], [rval1_scal])
                    ).make_node(xsym)
                if y < 0:
                    rval = [reciprocal(rval1)]
                else:
                    rval = [rval1]
            if rval:
                rval[0] = cast(rval[0], odtype)
                assert rval[0].type == node.outputs[0].type, (rval, node.outputs)
                return rval


@register_specialize
@local_optimizer([mul])
def local_mul_specialize(fgraph, node):
    """
    Remove special-case constants from mul arguments and useless neg in inputs.

    mul(-1, x) -> neg(x)
    mul(1, x, y) -> mul(x, y)
    mul(0, ...) -> alloc(0, shapes...)

    This is not done if we would add more nodes in the graph, like with:

    mul(-1, x, y) -/-> neg(mul(x, y))

    """
    # here, we are past the point of canonicalization, so we don't
    # want to put in un-necessary fills.
    #
    # at this point [post canonicalize], mul() may have many inputs.
    if node.op == mul:
        # the idea here is that we have pow(x, y)
        has_neg = False
        new_inputs = []
        nb_neg_node = 0
        nb_cst = 0
        for inp in node.inputs:
            # remove any neg arguments
            while inp.owner and inp.owner.op == neg:
                has_neg ^= True
                inp = inp.owner.inputs[0]
                nb_neg_node += 1

            # remove special case arguments of 1, -1 or 0
            y = get_constant(inp)
            if y == 1.0:
                nb_cst += 1
            elif y == -1.0:
                nb_cst += 1
                has_neg ^= True  # toggles
            elif y == 0.0:
                # if we find any zero, we just return right away
                return [broadcast_like(0, node.outputs[0], fgraph)]
            else:
                new_inputs.append(inp)

        if new_inputs != node.inputs:
            if new_inputs:
                if len(new_inputs) == 1:
                    if has_neg:
                        if new_inputs[0].dtype in (uint_dtypes + ["bool"]):
                            return
                        else:
                            rval = -new_inputs[0]
                    else:
                        rval = new_inputs[0]
                else:
                    # The next case would cause a replace by an equivalent case.
                    if has_neg and nb_neg_node == 0 and nb_cst == 1:
                        return
                    elif has_neg:
                        # Don't add an extra neg node as we can't
                        # fully replace this mul by a neg.
                        m1 = np.asarray(-1, dtype=node.outputs[0].dtype)
                        new_inputs = [m1] + new_inputs
                    rval = mul(*new_inputs)

                return [broadcast_like(rval, node.outputs[0], fgraph)]
            else:
                # there are no variable inputs to mul
                # N.B. this could have been constant-folded...
                if has_neg:
                    return [broadcast_like(-1, node.outputs[0], fgraph)]
                else:
                    return [broadcast_like(1, node.outputs[0], fgraph)]


@register_specialize
@local_optimizer([add])
def local_add_specialize(fgraph, node):
    """Remove zeros from ``add``s.

    TODO: This should be a canonicalization, no?
    """
    # here, we are past the point of canonicalization, so we don't want
    # to put in un-necessary fills.
    if node.op != add:
        return False

    new_inputs = []
    for inp in node.inputs:
        try:
            y = get_scalar_constant_value(inp)
        except NotScalarConstantError:
            y = inp
        if np.all(y == 0.0):
            continue
        new_inputs.append(inp)

    if len(new_inputs) == len(node.inputs):
        return False

    node_output = node.outputs[0]
    dtype = node_output.type.dtype

    if len(new_inputs) == 0:
        # we got rid of the entire expression!
        ndim = node_output.type.ndim
        # Reuse call to constant for cache()
        cst = constant(np.zeros((1,) * ndim, dtype=dtype))
        assert cst.type.broadcastable == (True,) * ndim
        return fill_chain(cst, node.inputs)

    if len(new_inputs) == 1:
        ret = fill_chain(new_inputs[0], node.inputs)
    else:
        ret = fill_chain(add(*new_inputs), node.inputs)

    # The dtype should not be changed. It can happen if the input
    # that was forcing upcasting was equal to 0.
    if ret[0].dtype != dtype:
        ret = [cast(ret[0], dtype)]

    return ret


mul_canonizer = in2out(
    LocalOptGroup(local_mul_canonizer, local_fill_sink, apply_all_opts=True),
    name="mul_canonizer_groups",
)


def check_for_x_over_absX(numerators, denominators):
    """Convert x/abs(x) into sign(x). """
    # TODO: this function should dig/search through dimshuffles
    # This won't catch a dimshuffled absolute value
    for den in list(denominators):
        if den.owner and den.owner.op == at_abs and den.owner.inputs[0] in numerators:
            if den.owner.inputs[0].type.dtype.startswith("complex"):
                # TODO: Make an Op that projects a complex number to
                #      have unit length but projects 0 to 0.  That
                #      would be a weird Op, but consistent with the
                #      special case below.  I heard there's some
                #      convention in Matlab that is similar to
                #      this... but not sure.
                pass
            else:
                denominators.remove(den)
                numerators.remove(den.owner.inputs[0])
                numerators.append(sgn(den.owner.inputs[0]))
    return numerators, denominators


local_mul_canonizer.add_simplifier(check_for_x_over_absX, "X_over_absX")


@register_canonicalize
@local_optimizer([at_abs])
def local_abs_lift(fgraph, node):
    """
    Move the abs toward the input.

    This is needed for check_for_x_over_absX to apply in more case.

    """
    if node.op == at_abs and node.inputs[0].owner:
        assert node.nin == 1
        if node.inputs[0].owner.op == mul:
            return [mul(*[at_abs(i) for i in node.inputs[0].owner.inputs])]
        if node.inputs[0].owner.op == true_div:
            i = node.inputs[0].owner.inputs
            return [true_div(at_abs(i[0]), at_abs(i[1]))]


@register_specialize
@local_optimizer([mul, true_div])
def local_abs_merge(fgraph, node):
    """
    Merge abs generated by local_abs_lift when the canonizer don't
    need it anymore

    """
    if (
        node.op == mul
        and sum([i.owner.op == at_abs for i in node.inputs if i.owner]) > 1
    ):
        inputs = []
        for i in node.inputs:
            if i.owner and i.owner.op == at_abs:
                inputs.append(i.owner.inputs[0])
            elif isinstance(i, Constant):
                try:
                    const = get_scalar_constant_value(i, only_process_constants=True)
                except NotScalarConstantError:
                    return False
                if not (const >= 0).all():
                    return False
                inputs.append(i)
            else:
                return False
        return [at_abs(mul(*inputs))]
    if (
        node.op == true_div
        and sum([i.owner.op == at_abs for i in node.inputs if i.owner]) == 2
    ):
        return [
            at_abs(
                true_div(node.inputs[0].owner.inputs[0], node.inputs[1].owner.inputs[0])
            )
        ]


@register_stabilize
@register_specialize
@local_optimizer([log])
def local_log1p(fgraph, node):
    # log(1+x) -> log1p(x)
    # log(1-x) -> log1p(-x)
    if node.op == log:
        (log_arg,) = node.inputs
        if log_arg.owner and log_arg.owner.op == add:
            scalars, scalar_inputs, nonconsts = scalarconsts_rest(
                log_arg.owner.inputs, only_process_constants=True
            )
            # scalar_inputs are potentially dimshuffled and fill'd scalars
            if scalars and np.allclose(np.sum(scalars), 1):
                if nonconsts:
                    if len(nonconsts) > 1:
                        ninp = add(*nonconsts)
                    else:
                        ninp = nonconsts[0]
                    if ninp.dtype != log_arg.type.dtype:
                        ninp = ninp.astype(node.outputs[0].dtype)
                    return fill_chain(log1p(ninp), scalar_inputs)

        elif log_arg.owner and log_arg.owner.op == sub:
            one = extract_constant(log_arg.owner.inputs[0], only_process_constants=True)
            if one != 1:
                return
            other = log_arg.owner.inputs[1]
            if other.dtype != log_arg.dtype:
                other = other.astype(log_arg.dtype)
            return [log1p(neg(other))]


@register_stabilize
@register_specialize
@local_optimizer([log])
def local_log_add_exp(fgraph, node):
    """
    ``log(exp(x)+exp(y)+exp(z)) = max + log(x-max, y-max, z-max)``

    TODO: in canonicalize, change log10 and log2 -> log
    """

    if node.op == log:
        z = node.inputs[0]
        if z.owner and z.owner.op == add:
            zi = z.owner.inputs
            pre_exp = [x.owner.inputs[0] for x in zi if x.owner and x.owner.op == exp]
            # all arguments to add are exp(<something>)
            if len(pre_exp) == len(zi):
                # Do not offset when max_pre = -np.inf, to avoid nan in the output
                # Switch statement is placed directly inside add to break the self-symmetry
                # of the returned output (otherwise the optimization would not stabilize)
                max_pre = reduce(maximum, pre_exp)
                ret = max_pre + log(
                    add(
                        *[
                            switch(isinf(max_pre), exp(max_pre), exp(p - max_pre))
                            for p in pre_exp
                        ]
                    )
                )
                return [ret]


@register_stabilize
@register_specialize
@local_optimizer([log])
def local_log_sum_exp(fgraph, node):
    # log(sum_i(exp(x_i))) = x_max + log(sum_i(exp(x_i - x_max)))

    if node.op != log:
        return

    sum_node = node.inputs[0].owner
    # If the sum has keepdims=True, there might be a dimshuffle
    if sum_node and isinstance(sum_node.op, DimShuffle):
        dimshuffle_op = sum_node.op
        sum_node = sum_node.inputs[0].owner
    else:
        dimshuffle_op = None

    if not sum_node or not isinstance(sum_node.op, Sum):
        return

    exp_node, axis = sum_node.inputs[0].owner, sum_node.op.axis
    if not exp_node or not (
        isinstance(exp_node.op, Elemwise) and isinstance(exp_node.op.scalar_op, aes.Exp)
    ):
        return

    pre_exp = exp_node.inputs[0]
    max_pre_exp = at_max(pre_exp, axis=axis)
    max_pre_exp_keepdims = makeKeepDims(pre_exp, max_pre_exp, axis)

    # Do not offset when max_pre = -np.inf, to avoid nan in the output
    # Switch statement is placed directly inside sum to break the self-symmetry
    # of the returned output (otherwise the optimization would not stabilize)
    ret = max_pre_exp + log(
        at_sum(
            switch(
                isinf(max_pre_exp_keepdims),
                exp(max_pre_exp_keepdims),
                exp(pre_exp - max_pre_exp_keepdims),
            ),
            axis=axis,
        ),
    )

    # Restore the dimshuffle op, if any.
    if dimshuffle_op:
        ret = dimshuffle_op(ret)

    return [ret]


def add_calculate(num, denum, aslist=False, out_type=None):
    # TODO: make sure that this function and mul_calculate are similar
    if out_type is None:
        zero = 0.0
    else:
        zero = _asarray(0, dtype=out_type.dtype)
    # zero = 0.0 if out_type is None else _asarray(0,
    # dtype=out_type.dtype)
    if out_type and out_type.dtype == "bool":
        if len(denum) == 0:
            # NumPy 1.14 do not accept to do "bool - bool"
            v = reduce(np.add, num, zero)
        else:
            raise Exception(
                "bool subtraction not supported. This should not happen as"
                " an earlier error should have been raised"
            )
    else:
        v = reduce(np.add, num, zero) - reduce(np.add, denum, zero)
    if aslist:
        if np.all(v == 0):
            return []
        else:
            return [v]
    return v


local_add_canonizer = AlgebraicCanonizer(add, sub, neg, add_calculate)
add_canonizer = in2out(
    LocalOptGroup(local_add_canonizer, local_fill_sink, apply_all_opts=True),
    name="add_canonizer_group",
)


register_canonicalize(local_add_canonizer, name="local_add_canonizer")


def distribute_greedy(pos_pairs, neg_pairs, num, denum, out_type, minscore=0):
    # each pair in pos_pairs and neg_pairs is a num/denum pair. this
    # function attempts to add num and denum to the corresponding parts
    # of each pair, and counts how many multiplications/divisions can
    # be saved in that way.

    # each division is counted like div_cost multiplications
    # (typically, division costs more so we are willing to multiply more
    # in order to divide less)
    # 1.5 was obtained through an informal test and may very well be
    # platform dependent
    div_cost = 1.5

    # score is number of operations saved, higher is better
    score = len(num) + div_cost * len(denum)
    new_pos_pairs = list(
        itertools.starmap(
            local_mul_canonizer.simplify,
            [(n + num, d + denum, out_type) for (n, d) in pos_pairs],
        )
    )
    new_neg_pairs = list(
        itertools.starmap(
            local_mul_canonizer.simplify,
            [(n + num, d + denum, out_type) for (n, d) in neg_pairs],
        )
    )
    for (n, d), (nn, dd) in zip(pos_pairs + neg_pairs, new_pos_pairs + new_neg_pairs):
        # We calculate how many operations we are saving with the new
        # num and denum
        score += len(n) + div_cost * len(d) - len(nn) - div_cost * len(dd)
    if score <= minscore:
        # the change is not applied because it adds too many operations
        return False, pos_pairs, neg_pairs
    return True, new_pos_pairs, new_neg_pairs


def attempt_distribution(factor, num, denum, out_type):
    """Try to insert each `num` and each `denum` in the factor?

    Returns
    -------
    changes?, new_factor, new_num, new_denum
        If there are changes, `new_num` and `new_denum` contain all the
        numerators and denominators that could not be distributed in the factor

    """
    pos_terms, neg_terms = local_add_canonizer.get_num_denum(factor)
    if len(pos_terms) == 1 and not neg_terms:
        return False, factor, num, denum
    pos_pairs = list(map(local_mul_canonizer.get_num_denum, pos_terms))
    neg_pairs = list(map(local_mul_canonizer.get_num_denum, neg_terms))
    change = False
    for n in list(num):
        success, pos_pairs, neg_pairs = distribute_greedy(
            pos_pairs, neg_pairs, [n], [], out_type
        )
        if success:
            change = True
            num.remove(n)
    for d in list(denum):
        success, pos_pairs, neg_pairs = distribute_greedy(
            pos_pairs, neg_pairs, [], [d], out_type
        )
        if success:
            change = True
            denum.remove(d)
    if not change:
        return change, factor, num, denum
    else:
        return (
            change,
            local_add_canonizer.merge_num_denum(
                list(itertools.starmap(local_mul_canonizer.merge_num_denum, pos_pairs)),
                list(itertools.starmap(local_mul_canonizer.merge_num_denum, neg_pairs)),
            ),
            num,
            denum,
        )


@register_canonicalize
@register_stabilize
@local_optimizer([mul, true_div, reciprocal])
def local_greedy_distributor(fgraph, node):
    """
    Optimize by reducing the number of multiplications and/or divisions.

    This optimization tries to apply distributivity of multiplication
    to addition in order to reduce the number of multiplications
    and/or divisions that must be done. The algorithm weighs division
    more than multiplication to account for the former's slightly
    greater computational cost.

    The following expressions are simplified:
    1. ((a/x + b/y) * x * y) --> a*y + b*x
    2. ((a/x + b) * x) --> a + b*x
    3. There are other forms too where node is a true_div.

    The following expressions are not simplified:
    4. ((a + b) * x) -/-> a*x + b*x

    This optimization aims to reduce computational cost. It may also
    increase numerical stability, e.g. when x and/or y tend to 0 in
    example 1.

    """

    out = node.outputs[0]
    num, denum = local_mul_canonizer.get_num_denum(out)
    if len(num) == 1 and not denum:
        return False

    new_num, new_denum = [], []

    change = False

    out_type = out.type
    for candidate in list(num):
        if candidate not in num:
            continue
        num.remove(candidate)
        _change, candidate, num, denum = attempt_distribution(
            candidate,
            num,
            denum,
            out_type,
        )

        change |= _change
        new_num.append(candidate)

    for candidate in list(denum):
        if candidate not in denum:
            continue
        denum.remove(candidate)
        _change, candidate, denum, num = attempt_distribution(
            candidate, denum, num, out_type
        )
        change |= _change
        new_denum.append(candidate)
    if not change:
        return False

    new_num += num
    new_denum += denum

    rval = local_mul_canonizer.merge_num_denum(new_num, new_denum)

    if rval.type != out.type:
        # WHY DOES THIS HAPPEN?
        return False

    return [rval]


get_clients_at_depth1 = partial(get_clients_at_depth, depth=1)
get_clients_at_depth2 = partial(get_clients_at_depth, depth=2)

# 1+erf(x)=>erfc(-x)
local_one_plus_erf = PatternSub(
    (add, 1, (erf, "x")),
    (erfc, (neg, "x")),
    allow_multiple_clients=True,
    name="local_one_plus_erf",
    tracks=[erf],
    get_nodes=get_clients_at_depth1,
)
register_canonicalize(local_one_plus_erf)
register_stabilize(local_one_plus_erf)
register_specialize(local_one_plus_erf)

# Only one of the two rewrites below is needed if a canonicalization is added
# for sub(x, y) -> add(x, -y) or a specialization for add(x, -y) -> sub(x, y)
# 1-erf(x)=>erfc(x)
local_one_minus_erf = PatternSub(
    (sub, 1, (erf, "x")),
    (erfc, "x"),
    allow_multiple_clients=True,
    name="local_one_minus_erf",
    tracks=[erf],
    get_nodes=get_clients_at_depth1,
)
register_canonicalize(local_one_minus_erf)
register_stabilize(local_one_minus_erf)
register_specialize(local_one_minus_erf)

local_one_minus_erf2 = PatternSub(
    (add, 1, (neg, (erf, "x"))),
    (erfc, "x"),
    allow_multiple_clients=True,
    name="local_one_minus_erf2",
    tracks=[erf],
    get_nodes=get_clients_at_depth2,
)
register_canonicalize(local_one_minus_erf2)
register_stabilize(local_one_minus_erf2)
register_specialize(local_one_minus_erf2)

# (-1)+erf(x) => -erfc(x)
# There is no need for erf(x)+(-1) nor erf(x) - 1, as the canonicalize will
# convert those to the matched pattern
local_erf_minus_one = PatternSub(
    (add, -1, (erf, "x")),
    (neg, (erfc, "x")),
    allow_multiple_clients=True,
    name="local_erf_minus_one",
    tracks=[erf],
    get_nodes=get_clients_at_depth1,
)
register_canonicalize(local_erf_minus_one)
register_stabilize(local_erf_minus_one)
register_specialize(local_erf_minus_one)

# Only one of the two rewrites below is needed if a canonicalization is added
# for sub(x, y) -> add(x, -y) or a specialization for add(x, -y) -> sub(x, y)
# 1-erfc(x) => erf(x)
local_one_minus_erfc = PatternSub(
    (sub, 1, (erfc, "x")),
    (erf, "x"),
    allow_multiple_clients=True,
    name="local_one_minus_erfc",
    tracks=[erfc],
    get_nodes=get_clients_at_depth1,
)
register_canonicalize(local_one_minus_erfc)
register_stabilize(local_one_minus_erfc)
register_specialize(local_one_minus_erfc)

local_one_minus_erfc2 = PatternSub(
    (add, 1, (neg, (erfc, "x"))),
    (erf, "x"),
    allow_multiple_clients=True,
    name="local_one_minus_erfc2",
    tracks=[erfc],
    get_nodes=get_clients_at_depth2,
)
register_canonicalize(local_one_minus_erfc2)
register_stabilize(local_one_minus_erfc2)
register_specialize(local_one_minus_erfc2)

# (-1)+erfc(-x)=>erf(x)
local_erf_neg_minus_one = PatternSub(
    (add, -1, (erfc, (neg, "x"))),
    (erf, "x"),
    allow_multiple_clients=True,
    name="local_erf_neg_minus_one",
    tracks=[erfc],
    get_nodes=get_clients_at_depth1,
)
register_canonicalize(local_erf_neg_minus_one)
register_stabilize(local_erf_neg_minus_one)
register_specialize(local_erf_neg_minus_one)


@register_stabilize
@register_specialize
@local_optimizer([log])
def local_log_erfc(fgraph, node):
    """Stability optimization for `log(erfc(x))`.

    log(erfc(x)) => when x>threshold,
                -x**2-log(x)-.5*log(pi)+log(1-1/(2*x**2)+3/(4*x**4)-15/(8*x**6))
    for float64: threshold=26.641747557 was chosen with:
    [(i,numpy.log(scipy.special.erfc(numpy.asarray([i],dtype='float64'))))
    for i in numpy.arange(26.641747557,26.6417475571,.00000000001)]
    for float32: threshold=10.0541949, [(i,numpy.log(scipy.special.erfc(
        numpy.asarray([i],dtype='float32')))) for i in numpy.arange(
        10.0541948,10.0541951,.0000001)]
    """
    if node.op != log:
        return False
    if not node.inputs[0].owner or node.inputs[0].owner.op != erfc:
        return False

    if hasattr(node.tag, "local_log_erfc_applied"):
        # We use that flag to don't apply the optimization recursively
        return False
    node.tag.local_log_erfc_applied = True

    x = node.inputs[0].owner.inputs[0]
    stab_value = (
        -(x ** 2)
        - log(x)
        - 0.5 * log(np.pi)
        + log(1 - 1 / (2 * x ** 2) + 3 / (4 * x ** 4) - 15 / (8 * x ** 6))
    )

    if node.outputs[0].dtype == "float32" or node.outputs[0].dtype == "float16":
        threshold = 10.0541949
    elif node.outputs[0].dtype == "float64":
        threshold = 26.641747557

    ret = switch(x < threshold, node.outputs[0], stab_value)
    ret.tag.values_eq_approx = values_eq_approx_remove_inf
    return [ret]


@register_stabilize
@register_specialize
@local_optimizer([true_div])
def local_grad_log_erfc_neg(fgraph, node):
    """Stability optimization for the grad of `log(erfc(x))`.

    ([y*]exp(-(x**2)))/erfc(x) # The y* is optional
    ([y*]exp(x**2))/erfc(-x) => [y*](when x > threshold,
                            sqrt(pi)*-x/(1-1/(2*x**2)+3/(4*x**4)-15/(8*x**6)))

    for float64: threshold=26.63 see at the end of the fct for the explanation
    for float32: threshold=9.3 see at the end of the fct for the explanation

    TODO: remove the constraint that there are only 2 inputs to exp(x**2)
        is the second.
    TODO: at the test point 10 in float32, there is instability in the original
        value. The original gives -30.0, the stab -20.1 and in float64 -18.1.
        Make it so that the test does not generate an error in that case!

    """
    if node.op != true_div:
        return False
    if not node.inputs[1].owner or node.inputs[1].owner.op != erfc:
        return False

    erfc_in = node.inputs[1]
    erfc_x = erfc_in.owner.inputs[0]

    if not node.inputs[0].owner:
        return False

    # TODO: All of this should be replaced with a single, simple unification
    # The mul is optional.
    if node.inputs[0].owner.op != mul:
        mul_in = None
        y = []
        if not node.inputs[0].owner or node.inputs[0].owner.op != exp:
            return False
        exp_in = node.inputs[0]
    else:
        mul_in = node.inputs[0]
        exp_in = None
        for idx, inp in enumerate(mul_in.owner.inputs):
            if inp.owner and inp.owner.op == exp:
                exp_in = inp
                break
        else:
            return False

        if len(mul_in.owner.inputs) == 2:
            y = [mul_in.owner.inputs[1 - idx]]
        else:
            y = mul_in.owner.inputs[:]
            del y[idx]

    if not exp_in.owner.inputs[0].owner:
        return False

    if exp_in.owner.inputs[0].owner.op == neg:
        neg_in = exp_in.owner.inputs[0]
        if not neg_in.owner.inputs[0].owner or neg_in.owner.inputs[0].owner.op != sqr:
            return False
        sqr_in = neg_in.owner.inputs[0]
        x = sqr_in.owner.inputs[0]
    elif exp_in.owner.inputs[0].owner.op == mul:
        # We should compare that -(erfc_x**2) is equivalent to mul_neg.
        # There is currently no easy way to do this in the general case,
        # so we implement some common case for now.

        # In many cases the neg are replaced by mul in the graph.
        # This also allows to stabilize log(erfc(cst*x)).
        mul_neg = exp_in.owner.inputs[0]

        # In case that multiple mul are not fused together, we do it here.
        def check_input(inputs):
            new_inputs = []
            for i in inputs:
                if i.owner and i.owner.op == mul:
                    new_inputs.extend(check_input(i.owner.inputs))
                else:
                    new_inputs.append(i)
            return new_inputs

        mul_inputs = check_input(mul_neg.owner.inputs)

        # Put the constant first.
        for i in range(len(mul_inputs)):
            if isinstance(i, Constant):
                if i == 0:
                    break
                else:
                    tmp = mul_inputs[0]
                    mul_inputs[0] = mul_inputs[i]
                    mul_inputs[i] = tmp
                    break
        mul_neg = mul(*mul_inputs)

        try:
            cst2 = get_scalar_constant_value(
                mul_neg.owner.inputs[0], only_process_constants=True
            )
        except NotScalarConstantError:
            return False

        if len(mul_neg.owner.inputs) == 2:
            if (
                not mul_neg.owner.inputs[1].owner
                or mul_neg.owner.inputs[1].owner.op != sqr
            ):
                return False
            sqr_in = mul_neg.owner.inputs[1]
            x = sqr_in.owner.inputs[0]
        elif len(mul_neg.owner.inputs) == 3:
            if mul_neg.owner.inputs[1] is not mul_neg.owner.inputs[2]:
                return False
            x = mul_neg.owner.inputs[1]
        else:
            return False

        if cst2 != -1:
            if (
                not erfc_x.owner
                or erfc_x.owner.op != mul
                or len(erfc_x.owner.inputs) != 2
            ):
                # todo implement that case
                return False
            if erfc_x.owner.inputs[1] is not mul_neg.owner.inputs[1]:
                return False

            x = erfc_x
            try:
                cst = get_scalar_constant_value(
                    erfc_x.owner.inputs[0], only_process_constants=True
                )
            except NotScalarConstantError:
                return False
            if cst2 != -cst * 2:
                return False

            # The constant is valid. Must check that the
        elif erfc_x is not x:
            return False

    else:
        return False

    if hasattr(node.tag, "local_grad_log_erfc_neg"):
        # We use that flag to don't apply the optimization recursively
        return False

    if erfc_x is not x:
        return None

    # we move the y outside the div.
    true_div_no_mul = true_div(exp_in, erfc_in)
    true_div_no_mul.owner.tag.local_grad_log_erfc_neg = True

    # aaron value
    stab_value = (
        x
        * at_pow(1 - 1 / (2 * (x ** 2)) + 3 / (4 * (x ** 4)) - 15 / (8 * (x ** 6)), -1)
        * cast(sqrt(np.pi), dtype=x.dtype)
    )

    if x.dtype == "float32" or x.dtype == "float16":
        threshold = 9.3
        # threshold = 10.1
    elif x.dtype == "float64":
        threshold = 26.641747557

    ret = switch(x < threshold, true_div_no_mul, stab_value)

    if y:
        ret = mul(ret, *y)

    ret.tag.values_eq_approx = values_eq_approx_remove_inf_nan

    return [ret]


def local_add_mul_fusion(fgraph, node):
    """Fuse consecutive add or mul in one such node with more inputs.

    It is better to fuse add/mul that way then in a Composite node as
    this make the inner graph of the Composite smaller. This allow to
    put more computation in a Composite before hitting the max
    recursion limit when pickling Composite.

    """
    if not isinstance(node.op, Elemwise) or not isinstance(
        node.op.scalar_op, (aes.Add, aes.Mul)
    ):
        return False

    s_op = node.op.scalar_op.__class__
    new_inp = []
    fused = False
    nb_inputs = len(node.inputs)
    max_inputs = float("inf")
    if hasattr(node.op, "max_inputs"):
        max_inputs = node.op.max_inputs(node)
    for inp in node.inputs:
        if (
            inp.owner
            and isinstance(inp.owner.op, Elemwise)
            and isinstance(inp.owner.op.scalar_op, s_op)
            and
            # Do not duplicate the operation.
            len(fgraph.clients[inp]) == 1
            and (nb_inputs + len(inp.owner.inputs) - 1) <= max_inputs
        ):
            new_inp.extend(inp.owner.inputs)
            fused = True
        else:
            new_inp.append(inp)

    # We can not compare the number of inputs as Mul and Add could have
    # 0 or 1 inputs in some corner cases.
    if fused:
        output = node.op(*new_inp)
        copy_stack_trace(node.outputs[0], output)

        # Do the recursion here to help lower the number of
        # FusionOptimizer iteration.
        if output.owner:
            output2 = local_add_mul_fusion(fgraph, output.owner)
            if output2:
                return output2
        return [output]


fuse_seqopt.register(
    "local_add_mul_fusion",
    FusionOptimizer(local_add_mul_fusion),
    "fast_run",
    "fusion",
    position=0,
)


def _skip_mul_1(r):
    if r.owner and r.owner.op == mul:
        not_is_1 = [i for i in r.owner.inputs if not _is_1(i)]
        if len(not_is_1) == 1:
            return not_is_1[0]


def _is_1(expr):
    """

    Returns
    -------
    bool
        True iff expr is a constant close to 1.

    """
    try:
        v = get_scalar_constant_value(expr)
        return np.allclose(v, 1)
    except NotScalarConstantError:
        return False


logsigm_to_softplus = PatternSub(
    (log, (sigmoid, "x")),
    (neg, (softplus, (neg, "x"))),
    allow_multiple_clients=True,
    values_eq_approx=values_eq_approx_remove_inf,
    skip_identities_fn=_skip_mul_1,
    tracks=[sigmoid],
    get_nodes=get_clients_at_depth1,
)
log1msigm_to_softplus = PatternSub(
    (log, (sub, dict(pattern="y", constraint=_is_1), (sigmoid, "x"))),
    (neg, (softplus, "x")),
    allow_multiple_clients=True,
    values_eq_approx=values_eq_approx_remove_inf,
    skip_identities_fn=_skip_mul_1,
    tracks=[sigmoid],
    get_nodes=get_clients_at_depth2,
)
log1pexp_to_softplus = PatternSub(
    (log1p, (exp, "x")),
    (softplus, "x"),
    values_eq_approx=values_eq_approx_remove_inf,
    allow_multiple_clients=True,
)
log1p_neg_sigmoid = PatternSub(
    (log1p, (neg, (sigmoid, "x"))),
    (neg, (softplus, "x")),
    values_eq_approx=values_eq_approx_remove_inf,
    allow_multiple_clients=True,
    tracks=[sigmoid],
    get_nodes=get_clients_at_depth2,
)

register_stabilize(logsigm_to_softplus, name="logsigm_to_softplus")
register_stabilize(log1msigm_to_softplus, name="log1msigm_to_softplus")
register_stabilize(log1pexp_to_softplus, name="log1pexp_to_softplus")
register_stabilize(log1p_neg_sigmoid, name="log1p_neg_sigmoid")
register_specialize(log1p_neg_sigmoid, name="log1p_neg_sigmoid")


def is_1pexp(t, only_process_constants=True):
    """

    Returns
    -------
    object
        If 't' is of the form (1+exp(x)), return (False, x).
        Else return None.

    """
    if t.owner and t.owner.op == add:
        scalars, scalar_inputs, nonconsts = scalarconsts_rest(
            t.owner.inputs, only_process_constants=only_process_constants
        )
        # scalar_inputs are potentially dimshuffled and filled with scalars
        if len(nonconsts) == 1:
            maybe_exp = nonconsts[0]
            if maybe_exp.owner and maybe_exp.owner.op == exp:
                # Verify that the constant terms sum to 1.
                if scalars:
                    scal_sum = scalars[0]
                    for s in scalars[1:]:
                        scal_sum = scal_sum + s
                    if np.allclose(scal_sum, 1):
                        return False, maybe_exp.owner.inputs[0]
    return None


def is_exp(var):
    """
    Match a variable with either of the `exp(x)` or `-exp(x)` patterns.

    Parameters
    ----------
    var
        The Variable to analyze.

    Returns
    -------
    tuple
        A pair (b, x) with `b` a boolean set to True if `var` is of the
        form `-exp(x)` and False if `var` is of the form `exp(x)`. If `var`
        cannot be cast into either form, then return `None`.

    """
    _neg = False
    neg_info = is_neg(var)
    if neg_info is not None:
        _neg = True
        var = neg_info
    if var.owner and var.owner.op == exp:
        return _neg, var.owner.inputs[0]


def is_mul(var):
    """
    Match a variable with `x * y * z * ...`.

    Parameters
    ----------
    var
        The Variable to analyze.

    Returns
    -------
    object
        A list [x, y, z, ...] if `var` is of the form `x * y * z * ...`,
        or None if `var` cannot be cast into this form.

    """
    if var.owner and var.owner.op == mul:
        return var.owner.inputs
    else:
        return None


def partition_num_or_denom(r, f):
    if r.owner and r.owner.op == mul:
        a = r.owner.inputs
    else:
        a = [r]

    # ugly 2.4-compatible thing
    f_terms = []
    _neg = False
    rest = []
    for t in a:
        f_t = f(t)
        if f_t is None:
            rest.append(t)
        else:
            neg_t, f_t = f_t
            f_terms.append(f_t)
            _neg ^= neg_t  # bit flip if neg_t is true
    return f_terms, rest, _neg


def is_neg(var):
    """
    Match a variable with the `-x` pattern.

    Parameters
    ----------
    var
        The Variable to analyze.

    Returns
    -------
    object
        `x` if `var` is of the form `-x`, or None otherwise.

    """
    var_node = var.owner
    if not var_node:
        return None
    # First match against `neg`.
    if var_node.op == neg:
        return var_node.inputs[0]
    # Then match against a multiplication by -1.
    if var_node.op == mul and len(var_node.inputs) >= 2:
        for idx, mul_input in enumerate(var_node.inputs):
            try:
                constant = get_scalar_constant_value(mul_input)
                is_minus_1 = np.allclose(constant, -1)
            except NotScalarConstantError:
                is_minus_1 = False
            if is_minus_1:
                # Found a multiplication by -1.
                if len(var_node.inputs) == 2:
                    # Only return the other input.
                    return var_node.inputs[1 - idx]
                else:
                    # Return the multiplication of all other inputs.
                    return mul(*(var_node.inputs[0:idx] + var_node.inputs[idx + 1 :]))
    # No match.
    return None


@register_stabilize
@local_optimizer([true_div])
def local_exp_over_1_plus_exp(fgraph, node):
    """
    exp(x)/(1+exp(x)) -> sigm(x)
    c/(1+exp(x)) -> c*sigm(-x)

    """
    # this optimization should be done for numerical stability
    # so we don't care to check client counts
    if node.op == true_div:

        # find all the exp() terms in the numerator
        num, denom = node.inputs
        num_exp_x, num_rest, num_neg = partition_num_or_denom(num, is_exp)
        denom_1pexp, denom_rest, denom_neg = partition_num_or_denom(denom, is_1pexp)

        sigmoids = []
        for t in denom_1pexp:
            if t in num_exp_x:
                # case: exp(x) /(1+exp(x))
                sigmoids.append(sigmoid(t))
                del num_exp_x[num_exp_x.index(t)]
            else:
                # case: 1/(1+exp(x))
                sigmoids.append(sigmoid(-t))
            copy_stack_trace(node.outputs[0], sigmoids[-1])

        if not sigmoids:  # we didn't find any.  abort
            return
        # put the new numerator together
        new_num = sigmoids + [exp(t) for t in num_exp_x] + num_rest
        if len(new_num) == 1:
            new_num = new_num[0]
        else:
            new_num = mul(*new_num)

        if num_neg ^ denom_neg:
            new_num = -new_num

        copy_stack_trace(num, new_num)

        if len(denom_rest) == 0:
            return [new_num]
        elif len(denom_rest) == 1:
            out = new_num / denom_rest[0]
        else:
            out = new_num / mul(*denom_rest)

        copy_stack_trace(node.outputs[0], out)
        return [out]


def parse_mul_tree(root):
    """
    Parse a tree of multiplications starting at the given root.

    Parameters
    ----------
    root
        The variable at the root of the tree.

    Returns
    -------
    object
        A tree where each non-leaf node corresponds to a multiplication
        in the computation of `root`, represented by the list of its inputs.
        Each input is a pair [n, x] with `n` a boolean value indicating whether
        sub-tree `x` should be negated.

    Examples
    --------

    .. code-block:: python

        x * y               -> [False, [[False, x], [False, y]]]
        -(x * y)            -> [True, [[False, x], [False, y]]]
        -x * y              -> [False, [[True, x], [False, y]]]
        -x                  -> [True, x]
        (x * y) * -z        -> [False, [[False, [[False, x], [False, y]]],
                                        [True, z]]]

    """
    # Is it a multiplication?
    mul_info = is_mul(root)
    if mul_info is None:
        # Is it a negation?
        neg_info = is_neg(root)
        if neg_info is None:
            # Keep the root "as is".
            return [False, root]
        else:
            # Recurse, inverting the negation.
            neg, sub_tree = parse_mul_tree(neg_info)
            return [not neg, sub_tree]
    else:
        # Recurse into inputs.
        return [False, list(map(parse_mul_tree, mul_info))]


def replace_leaf(arg, leaves, new_leaves, op, neg):
    """
    Attempt to replace a leaf of a multiplication tree.

    We search for a leaf in `leaves` whose argument is `arg`, and if we find
    one, we remove it from `leaves` and add to `new_leaves` a leaf with
    argument `arg` and variable `op(arg)`.

    Parameters
    ----------
    arg
        The argument of the leaf we are looking for.
    leaves
        List of leaves to look into. Each leaf should be a pair
        (x, l) with `x` the argument of the Op found in the leaf, and `l` the
        actual leaf as found in a multiplication tree output by `parse_mul_tree`
        (i.e. a pair [boolean, variable]).
    new_leaves
        If a replacement occurred, then the leaf is removed from `leaves`
        and added to the list `new_leaves` (after being modified by `op`).
    op
        A function that, when applied to `arg`, returns the Variable
        we want to replace the original leaf variable with.
    neg : bool
        If True, then the boolean value associated to the leaf should
        be swapped. If False, then this value should remain unchanged.

    Returns
    -------
    bool
        True if a replacement occurred, or False otherwise.

    """
    for idx, x in enumerate(leaves):
        if x[0] == arg:
            x[1][0] ^= neg
            x[1][1] = op(arg)
            leaves.pop(idx)
            new_leaves.append(x)
            return True
    return False


def simplify_mul(tree):
    """
    Simplify a multiplication tree.

    Parameters
    ----------
    tree
        A multiplication tree (as output by `parse_mul_tree`).

    Returns
    -------
    object
        A multiplication tree computing the same output as `tree` but without
        useless multiplications by 1 nor -1 (identified by leaves of the form
        [False, None] or [True, None] respectively). Useless multiplications
        (with less than two inputs) are also removed from the tree.

    """
    neg, inputs = tree
    if isinstance(inputs, list):
        # Recurse through inputs.
        s_inputs = []
        for s_i in map(simplify_mul, inputs):
            if s_i[1] is None:
                # Multiplication by +/-1.
                neg ^= s_i[0]
            else:
                s_inputs.append(s_i)
        if not s_inputs:
            # The multiplication is empty.
            rval = [neg, None]
        elif len(s_inputs) == 1:
            # The multiplication has a single input.
            s_inputs[0][0] ^= neg
            rval = s_inputs[0]
        else:
            rval = [neg, s_inputs]
    else:
        rval = tree
    # print 'simplify_mul: %s -> %s' % (tree, rval)
    return rval


def compute_mul(tree):
    """
    Compute the Variable that is the output of a multiplication tree.

    This is the inverse of the operation performed by `parse_mul_tree`, i.e.
    compute_mul(parse_mul_tree(tree)) == tree.

    Parameters
    ----------
    tree
        A multiplication tree (as output by `parse_mul_tree`).

    Returns
    -------
    object
        A Variable that computes the multiplication represented by the tree.

    """
    neg, inputs = tree
    if inputs is None:
        raise AssertionError(
            "Function `compute_mul` found a missing leaf, did you forget to "
            "call `simplify_mul` on the tree first?"
        )
    elif isinstance(inputs, list):
        # Recurse through inputs.
        rval = mul(*list(map(compute_mul, inputs)))
    else:
        rval = inputs
    if neg:
        rval = -rval
    return rval


def perform_sigm_times_exp(
    tree,
    exp_x=None,
    exp_minus_x=None,
    sigm_x=None,
    sigm_minus_x=None,
    parent=None,
    child_idx=None,
    full_tree=None,
):
    """
    Core processing of the `local_sigm_times_exp` optimization.

    This recursive function operates on a multiplication tree as output by
    `parse_mul_tree`. It walks through the tree and modifies it in-place
    by replacing matching pairs (exp, sigmoid) with the desired optimized
    version.

    Parameters
    ----------
    tree
        The sub-tree to operate on.
    exp_x
        List of arguments x so that `exp(x)` exists somewhere in the whole
        multiplication tree. Each argument is a pair (x, leaf) with `x` the
        argument of the exponential, and `leaf` the corresponding leaf in the
        multiplication tree (of the form [n, exp(x)] -- see `parse_mul_tree`).
        If None, this argument is initialized to an empty list.
    exp_minus_x
        Similar to `exp_x`, but for `exp(-x)`.
    sigm_x
        Similar to `exp_x`, but for `sigmoid(x)`.
    sigm_minus_x
        Similar to `exp_x`, but for `sigmoid(-x)`.
    parent
        Parent of `tree` (None if `tree` is the global root).
    child_idx
        Index of `tree` in its parent's inputs (None if `tree` is the global
        root).
    full_tree
        The global multiplication tree (should not be set except by recursive
        calls to this function). Used for debugging only.

    Returns
    -------
    bool
        True if a modification was performed somewhere in the whole multiplication
        tree, or False otherwise.

    """
    if exp_x is None:
        exp_x = []
    if exp_minus_x is None:
        exp_minus_x = []
    if sigm_x is None:
        sigm_x = []
    if sigm_minus_x is None:
        sigm_minus_x = []
    if full_tree is None:
        full_tree = tree
    if False:  # Debug code.
        print("<perform_sigm_times_exp>")
        print(f"  full_tree   = {full_tree}")
        print(f"  tree        = {tree}")
        print(f"  exp_x       = {exp_x}")
        print(f"  exp_minus_x = {exp_minus_x}")
        print(f"  sigm_x      = {sigm_x}")
        print(f"  sigm_minus_x= {sigm_minus_x}")
    neg, inputs = tree
    if isinstance(inputs, list):
        # Recurse through inputs of the multiplication.
        rval = False
        for sub_idx, sub_tree in enumerate(inputs):
            rval |= perform_sigm_times_exp(
                tree=sub_tree,
                parent=tree,
                child_idx=sub_idx,
                exp_x=exp_x,
                exp_minus_x=exp_minus_x,
                sigm_x=sigm_x,
                sigm_minus_x=sigm_minus_x,
                full_tree=full_tree,
            )
        return rval
    else:
        # Reached a leaf: if it is an exponential or a sigmoid, then we
        # first attempt to find a match in leaves already visited.
        # If there is such a match, we modify the already-visited leaf
        # accordingly: for instance if we visited a leaf sigmoid(x), then
        # find later a -exp(-x), we replace the previous leaf by
        # -sigmoid(-x) and remove the -exp(-x) from the tree.
        # If no match is found, then we register this leaf so that it can
        # be found later while walking the tree.
        var = inputs
        keep_it = False
        exp_info = is_exp(var)
        if exp_info is not None:
            exp_neg, exp_arg = exp_info
            neg ^= exp_neg
            neg_arg = is_neg(exp_arg)
            if neg_arg is None:
                if not replace_leaf(exp_arg, sigm_minus_x, sigm_x, sigmoid, neg):
                    exp_x.append((exp_arg, tree))
                    keep_it = True
            else:
                if not replace_leaf(
                    neg_arg, sigm_x, sigm_minus_x, lambda x: sigmoid(-x), neg
                ):
                    exp_minus_x.append((neg_arg, tree))
                    keep_it = True
        elif var.owner and var.owner.op == sigmoid:
            sigm_arg = var.owner.inputs[0]
            neg_arg = is_neg(sigm_arg)
            if neg_arg is None:
                if not replace_leaf(
                    sigm_arg, exp_minus_x, sigm_minus_x, lambda x: sigmoid(-x), neg
                ):
                    sigm_x.append((sigm_arg, tree))
                    keep_it = True
            else:
                if not replace_leaf(neg_arg, exp_x, sigm_x, sigmoid, neg):
                    sigm_minus_x.append((neg_arg, tree))
                    keep_it = True
        else:
            # It is not an exponential nor a sigmoid.
            keep_it = True
        if not keep_it:
            # Delete this leaf, i.e. replace it by [False, None] (corresponding
            # to a multiplication by 1).
            assert parent is not None
            parent[1][child_idx] = [False, None]
        return not keep_it


@register_stabilize
@local_optimizer([mul])
def local_sigm_times_exp(fgraph, node):
    """
    exp(x) * sigm(-x) -> sigm(x)
    exp(-x) * sigm(x) -> sigm(-x)

    todo: add stack traces to the intermediate variables
    """
    # Bail early if it is not a multiplication.
    if node.op != mul:
        return None
    # Obtain tree of multiplications starting at this node.
    mul_tree = parse_mul_tree(node.outputs[0])
    # Perform core optimization.
    did_something = perform_sigm_times_exp(mul_tree)
    if not did_something:
        # No change.
        return None
    # The optimization may have introduced multiplications by 1 in the tree:
    # get rid of them.
    mul_tree = simplify_mul(mul_tree)
    # Recompute final output based on the updated tree.
    out = compute_mul(mul_tree)
    # keep the stack trace
    copy_stack_trace(node.outputs[0], out)
    return [out]


@register_stabilize
@local_optimizer([reciprocal])
def local_reciprocal_1_plus_exp(fgraph, node):
    """``reciprocal(1+exp(x)) -> sigm(-x)``

    TODO: This is redundant; we can just decided on *one* canonical form
    for division (e.g. either `true_div` or `reciprocal`) and have this
    taken care of with existing rewrites.
    """
    # this optimization should be done for numerical stability
    # so we don't care to check client counts
    if node.op == reciprocal:
        reciprocal_arg = node.inputs[0]
        if reciprocal_arg.owner and reciprocal_arg.owner.op == add:
            scalars_, scalar_inputs, nonconsts = scalarconsts_rest(
                reciprocal_arg.owner.inputs, only_process_constants=True
            )
            # scalar_inputs are potentially dimshuffled and fill'd scalars
            if len(nonconsts) == 1:
                if nonconsts[0].owner and nonconsts[0].owner.op == exp:
                    if scalars_ and np.allclose(np.sum(scalars_), 1):
                        out = fill_chain(
                            sigmoid(neg(nonconsts[0].owner.inputs[0])),
                            scalar_inputs,
                        )
                        # keep combined stack traces of
                        #     exp(x):           nonconsts[0],
                        #     1 + exp(x):       reciprocal_arg,
                        #     1 / (1 + exp(x)): node.outputs[0]
                        copy_stack_trace(
                            [nonconsts[0], reciprocal_arg, node.outputs[0]], out
                        )
                        return out


# 1 - sigmoid(x) -> sigmoid(-x)
local_1msigmoid = PatternSub(
    (sub, dict(pattern="y", constraint=_is_1), (sigmoid, "x")),
    (sigmoid, (neg, "x")),
    tracks=[sigmoid],
    get_nodes=get_clients_at_depth1,
    name="local_1msigmoid",
)
register_stabilize(local_1msigmoid)
register_specialize(local_1msigmoid)


log1pmexp_to_log1mexp = PatternSub(
    (log1p, (neg, (exp, "x"))),
    (log1mexp, "x"),
    allow_multiple_clients=True,
)
register_stabilize(log1pmexp_to_log1mexp, name="log1pmexp_to_log1mexp")


# log(sigmoid(x) / (1 - sigmoid(x))) -> x
# i.e logit(sigmoid(x)) -> x
local_logit_sigmoid = PatternSub(
    (log, (true_div, (sigmoid, "x"), (sub, 1, (sigmoid, "x")))),
    "x",
    tracks=[sigmoid],
    get_nodes=get_clients_at_depth2,
    allow_multiple_clients=True,
    name="local_logit_sigmoid",
)
register_canonicalize(local_logit_sigmoid)
register_specialize(local_logit_sigmoid)


# sigmoid(log(x / (1-x)) -> x
# i.e., sigmoid(logit(x)) -> x
local_sigmoid_logit = PatternSub(
    (sigmoid, (log, (true_div, "x", (sub, 1, "x")))),
    "x",
    allow_multiple_clients=True,
    name="local_sigmoid_logit",
)
register_canonicalize(local_sigmoid_logit)
register_specialize(local_sigmoid_logit)
