@register_canonicalize
@register_specialize
@gof.local_optimizer([T.Sum, T.elemwise.Prod])
def local_sum_prod_div_dimshuffle(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, (T.Sum, T.elemwise.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 == T.true_div:
            numerator, denominator = node_input.owner.inputs

            # Old, bugged logic, reproduced here only to warn users
            if (config.warn.sum_div_dimshuffle_bug and
                    isinstance(node.op, T.Sum) and
                    numerator.owner and
                    isinstance(numerator.owner.op, T.DimShuffle)):
                # Check compatibility
                new_order = numerator.owner.op.new_order
                compatible_dims = True
                for ax in axis:
                    if len(new_order) <= ax or new_order[ax] != 'x':
                        compatible_dims = False
                        break

                if compatible_dims:
                    _logger.warn('WARNING: Your current code is fine, but'
                                 ' Theano versions between '
                                 'rev. 3bd9b789f5e8 (2010-06-16) and'
                                 ' cfc6322e5ad4 (2010-08-03) would '
                                 'have given an incorrect result. '
                                 'To disable this warning, set the Theano'
                                 ' flag warn.sum_div_dimshuffle_bug to'
                                 ' False.')

            if denominator.owner and isinstance(denominator.owner.op,
                                                T.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 = T.DimShuffle(
                            dimshuffle_input.type.broadcastable,
                            optimized_dimshuffle_order)(dimshuffle_input)

                        if (config.warn.sum_div_dimshuffle_bug and
                                isinstance(node.op, T.Sum)):
                            _logger.warn('WARNING: Your current code is fine,'
                                         ' but Theano versions between '
                                         'rev. 3bd9b789f5e8 (2010-06-16) and'
                                         ' cfc6322e5ad4 (2010-08-03) would '
                                         'have given an incorrect result. '
                                         'To disable this warning, set the'
                                         ' Theano flag '
                                         'warn.sum_div_dimshuffle_bug'
                                         ' to False.')

                    if isinstance(node.op, T.Sum):
                        op_on_compatible_dims = T.sum(
                            numerator, axis=compatible_dims)
                        rval = T.true_div(
                            op_on_compatible_dims,
                            optimized_dimshuffle)
                        if len(reordered_incompatible_dims) > 0:
                            rval = T.sum(rval,
                                         axis=reordered_incompatible_dims)
                    elif isinstance(node.op, T.elemwise.Prod):
                        op_on_compatible_dims = T.prod(
                            numerator, axis=compatible_dims)
                        dtype = numerator.dtype
                        rval = T.true_div(
                            op_on_compatible_dims,
                            (optimized_dimshuffle **
                                T.prod([numerator.shape[ax].astype(dtype)
                                        for ax in compatible_dims])))
                        if len(reordered_incompatible_dims) > 0:
                            rval = T.prod(rval,
                                          axis=reordered_incompatible_dims)
                    return [rval]
dataset/ETHPy150Open rizar/attention-lvcsr/libs/Theano/theano/tensor/opt.py/local_sum_prod_div_dimshuffle/original