<commit_before>import collections

import numpy

from chainer import cuda
from chainer import function

_args = 'float* y, float* x, int cdimy, int cdimx, int rdim, int coffset'
_preamble = '''
#define COPY(statement) \
    int l   = i / (rdim * cdimy);  \
    int c   = i / rdim % cdimy + coffset;  \
    int r   = i % rdim;  \
    int idx = r + rdim * (c + cdimx * l);  \
    statement;
'''


class SplitAxis(function.Function):

    """Function that splits multiple arrays towards the specified axis."""

    def __init__(self, indices_or_sections, axis):
        if not isinstance(indices_or_sections, (int, collections.Iterable)):
            raise TypeError('indices_or_sections must be integer or 1-D array')
        self.indices_or_sections = indices_or_sections
        self.axis = axis

    def forward_cpu(self, x):
        if isinstance(self.indices_or_sections, collections.Iterable):
            cdimx = x[0].shape[self.axis]
            ind = list(self.indices_or_sections)
            ind.append(cdimx)
            prev_i = 0
            for i in ind:
                cdimy = max(0, min(i, cdimx) - prev_i)
                if cdimy == 0:
                    raise ValueError('Not support if shape contains 0')
                prev_i = i

        return tuple(numpy.split(x[0], self.indices_or_sections, self.axis))

    def forward_gpu(self, x):
        xshape = x[0].shape
        self.cdimx = xshape[self.axis]
        self.rdim = numpy.prod(xshape[self.axis + 1:])

        if isinstance(self.indices_or_sections, collections.Iterable):
            ind = list(self.indices_or_sections)
            ind.append(self.cdimx)
        else:
            sec = self.indices_or_sections
            if self.cdimx % sec:
                raise ValueError(
                    'array split does not result in an equal division')
            ind = numpy.arange(1, sec + 1) * (self.cdimx // sec)
        ys = []
        kernel = cuda.elementwise(
            _args, 'COPY(y[i] = x[idx])', 'split_fwd', preamble=_preamble)
        prev_i = 0
        for i in ind:
            cdimy = max(0, min(i, self.cdimx) - prev_i)
            s = list(xshape)
            s[self.axis] = cdimy
            y = cuda.empty(s, dtype=x[0].dtype)
            if cdimy == 0:
                raise ValueError('Not support if shape contains 0')
            kernel(y, x[0], cdimy, self.cdimx, self.rdim, prev_i)
            prev_i = i
            ys.append(y)
        return tuple(ys)

    def backward_cpu(self, x, gys):
        return numpy.concatenate(gys, axis=self.axis),

    def backward_gpu(self, x, gys):
        gx = cuda.empty_like(x[0])
        coffset = 0
        kernel = cuda.elementwise(
            _args, 'COPY(x[idx] = y[i])', 'split_bwd', preamble=_preamble)
        for gy in gys:
            cdimy = gy.shape[self.axis]
            if cdimy != 0:
                kernel(gy, gx, cdimy, self.cdimx, self.rdim, coffset)
            coffset += cdimy
        return gx,


def split_axis(x, indices_or_sections, axis):
    """Splits given variables along an axis.

    Args:
        x (tuple of Variables): Variables to be split.
        indices_or_sections (int or 1-D array): If this argument is an integer,
            N, the array will be divided into N equal arrays along axis.
            If it is a 1-D array of sorted integers, the entries
            indicate where along which axis the array is split.
        axis (int): Axis that the input arrays are split along.<commit_msg>Split minor function<commit_after>import collections

import numpy

from chainer import cuda
from chainer import function

_args = 'float* y, float* x, int cdimy, int cdimx, int rdim, int coffset'
_preamble = '''
#define COPY(statement) \
    int l   = i / (rdim * cdimy);  \
    int c   = i / rdim % cdimy + coffset;  \
    int r   = i % rdim;  \
    int idx = r + rdim * (c + cdimx * l);  \
    statement;
'''


class SplitAxis(function.Function):

    """Function that splits multiple arrays towards the specified axis."""

    def __init__(self, indices_or_sections, axis):
        if not isinstance(indices_or_sections, (int, collections.Iterable)):
            raise TypeError('indices_or_sections must be integer or 1-D array')
        self.indices_or_sections = indices_or_sections
        self.axis = axis

    def forward_cpu(self, x):
        return tuple(numpy.split(x[0], self.indices_or_sections, self.axis))

    def forward_gpu(self, x):
        xshape = x[0].shape
        self.cdimx = xshape[self.axis]
        self.rdim = numpy.prod(xshape[self.axis + 1:])

        if isinstance(self.indices_or_sections, collections.Iterable):
            ind = list(self.indices_or_sections)
            ind.append(self.cdimx)
        else:
            sec = self.indices_or_sections
            if self.cdimx % sec:
                raise ValueError(
                    'array split does not result in an equal division')
            ind = numpy.arange(1, sec + 1) * (self.cdimx // sec)
        ys = []
        kernel = cuda.elementwise(
            _args, 'COPY(y[i] = x[idx])', 'split_fwd', preamble=_preamble)
        prev_i = 0
        for i in ind:
            cdimy = max(0, min(i, self.cdimx) - prev_i)
            s = list(xshape)
            s[self.axis] = cdimy
            y = cuda.empty(s, dtype=x[0].dtype)
            if cdimy != 0:
                kernel(y, x[0], cdimy, self.cdimx, self.rdim, prev_i)
            prev_i = i
            ys.append(y)
        return tuple(ys)

    def backward_cpu(self, x, gys):
        return numpy.concatenate(gys, axis=self.axis),

    def backward_gpu(self, x, gys):
        gx = cuda.empty_like(x[0])
        coffset = 0
        kernel = cuda.elementwise(
            _args, 'COPY(x[idx] = y[i])', 'split_bwd', preamble=_preamble)
        for gy in gys:
            cdimy = gy.shape[self.axis]
            if cdimy != 0:
                kernel(gy, gx, cdimy, self.cdimx, self.rdim, coffset)
            coffset += cdimy
        return gx,


def split_axis(x, indices_or_sections, axis):
    """Splits given variables along an axis.

    Args:
        x (tuple of Variables): Variables to be split.
        indices_or_sections (int or 1-D array): If this argument is an integer,
            N, the array will be divided into N equal arrays along axis.
            If it is a 1-D array of sorted integers, the entries
            indicate where along which axis the array is split.
        axis (int): Axis that the input arrays are split along.
{'fraction_non_alphanumeric': 0.10480887792848335, 'fraction_numerical': 0.005856966707768188, 'mean_word_length': 2.015334572490706, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '6177672', 'n_tokens_mistral': 2124, 'n_tokens_neox': 2062, 'n_words': 603}