<commit_before>from chainer.backends import cuda
from chainer import function_node
from chainer import utils


class DiagEinSum(function_node.FunctionNode):

    def __init__(self, in_subs, out_sub, out_shape=None):
        self.in_subs = in_subs
        self.out_sub = out_sub
        self.out_shape = out_shape

    def check_type_forward(self, in_types):
        pass

    def forward(self, inputs):
        n_args = len(inputs)
        self.retain_inputs(tuple(range(n_args)))
        """
        if n_args >= 2:
            self.retain_inputs(tuple(range(n_args)))
        else:
            self.input_shape = inputs[0].shape
        self.n_args = n_args
        """

        xp = cuda.get_array_module(inputs[0])

        out_sub = self.out_sub

        diag_map = []
        ein_sub = []
        for i, s in enumerate(out_sub):
            i0 = out_sub.index(s)
            if i == i0:
                if s in self.in_subs:
                    ein_sub.append(s)
                else:
                    i0 = None
            diag_map.append(i0)

        subscript = '{}->{}'.format(
            self.in_subs,
            ''.join(ein_sub)
        )
        y = utils.force_array(xp.einsum(subscript, *inputs))

        for i, i0 in enumerate(diag_map):
            if i0 is None:
                # broadcast to new axis
                assert self.out_shape is not None, \
                    "Give out_shape to put new subscripts in the result"
                y = xp.broadcast_to(
                    xp.expand_dims(y, axis=i),
                    y.shape[:i] + (self.out_shape[i],) + y.shape[i:]
                )
            elif i0 != i:
                # make diagonal
                size = y.shape[i0]
                z = xp.zeros(
                    y.shape[:i] + (size,) + y.shape[i:],
                    dtype=y.dtype)
                indexer = (
                    (slice(None),) * i0
                    + (xp.arange(size),)
                    + (slice(None),) * (i - i0 - 1))
                z[indexer + (xp.arange(size), Ellipsis,)] = \
                    y[indexer + (Ellipsis,)]
                y = z
        return y,

    def backward(self, indices, grad_outputs):
        inputs = self.get_retained_inputs()
        """
        if self.n_args >= 2:
            inputs = self.get_retained_inputs()
        else:
            # TODO(kataoka): fix
            xp = cuda.get_array_module(grad_outputs[0])
            inputs = [xp.zeros(self.input_shape, dtype=grad_outputs[0].dtype)]
        """

        g, = grad_outputs

        fwd_in_subs = self.in_subs.split(',')
        fwd_out_sub = self.out_sub
        return tuple(
            inputs[i] * 0. +  # it seems a bug of gradient_check
            DiagEinSum(
                in_subs=','.join([<commit_msg>refactor: y.shape<commit_after>from chainer.backends import cuda
from chainer import function_node
from chainer import utils


class DiagEinSum(function_node.FunctionNode):

    def __init__(self, in_subs, out_sub, out_shape=None):
        self.in_subs = in_subs
        self.out_sub = out_sub
        self.out_shape = out_shape

    def check_type_forward(self, in_types):
        pass

    def forward(self, inputs):
        n_args = len(inputs)
        self.retain_inputs(tuple(range(n_args)))
        """
        if n_args >= 2:
            self.retain_inputs(tuple(range(n_args)))
        else:
            self.input_shape = inputs[0].shape
        self.n_args = n_args
        """

        xp = cuda.get_array_module(inputs[0])

        out_sub = self.out_sub

        diag_map = []
        ein_sub = []
        for i, s in enumerate(out_sub):
            i0 = out_sub.index(s)
            if i == i0:
                if s in self.in_subs:
                    ein_sub.append(s)
                else:
                    i0 = None
            diag_map.append(i0)

        subscript = '{}->{}'.format(
            self.in_subs,
            ''.join(ein_sub)
        )
        y = utils.force_array(xp.einsum(subscript, *inputs))

        shape = list(y.shape)
        for i, i0 in enumerate(diag_map):
            if i0 is None:
                # broadcast to new axis
                assert self.out_shape is not None, \
                    "Give out_shape to put new subscripts in the result"
                shape.insert(i, self.out_shape[i])
                y = xp.broadcast_to(xp.expand_dims(y, axis=i), shape)
            elif i0 != i:
                # make diagonal
                size = shape[i0]
                shape.insert(i, size)
                z = xp.zeros(shape, dtype=y.dtype)
                indexer = (
                    (slice(None),) * i0
                    + (xp.arange(size),)
                    + (slice(None),) * (i - i0 - 1))
                z[indexer + (xp.arange(size), Ellipsis,)] = \
                    y[indexer + (Ellipsis,)]
                y = z
        return y,

    def backward(self, indices, grad_outputs):
        inputs = self.get_retained_inputs()
        """
        if self.n_args >= 2:
            inputs = self.get_retained_inputs()
        else:
            # TODO(kataoka): fix
            xp = cuda.get_array_module(grad_outputs[0])
            inputs = [xp.zeros(self.input_shape, dtype=grad_outputs[0].dtype)]
        """

        g, = grad_outputs

        fwd_in_subs = self.in_subs.split(',')
        fwd_out_sub = self.out_sub
        return tuple(
            inputs[i] * 0. +  # it seems a bug of gradient_check
            DiagEinSum(
                in_subs=','.join([
{'fraction_non_alphanumeric': 0.11018635787950064, 'fraction_numerical': 0.006513479283517279, 'mean_word_length': 1.4341699691765741, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3606202', 'n_tokens_mistral': 1815, 'n_tokens_neox': 1733, 'n_words': 354}