<commit_before>import collections

import numpy
import six

from chainer import cuda
from chainer import function
from chainer.utils import type_check


_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 check_type_forward(self, in_types):
        type_check.expect(
            in_types.size() == 1,
            in_types[0].ndim >= self.axis
        )

        if isinstance(self.indices_or_sections, collections.Iterable):
            max_index = type_check.Variable(
                self.indices_or_sections[-1], 'max_index')
            type_check.expect(in_types[0].shape[self.axis] >= max_index)
        else:
            sections = type_check.Variable(
                self.indices_or_sections, 'sections')
            type_check.expect(in_types[0].shape[self.axis] % sections == 0)

    def check_type_backward(self, in_types, out_types):
        if isinstance(self.indices_or_sections, collections.Iterable):
            num_indices = type_check.Variable(
                len(self.indices_or_sections), 'len(indices)')
            output_size = num_indices + 1
        else:
            output_size = type_check.Variable(
                self.indices_or_sections, 'sections')

        type_check.expect(out_types.size() == output_size.eval())

        for out_type in out_types:
            type_check.expect(
                out_type.ndim == in_types[0].ndim,
                out_type.dtype == out_types[0].dtype
            )

        for out_type in out_types:
            for i in six.moves.range(in_types[0].ndim.eval()):
                if i == self.axis:
                    continue
                type_check.expect(out_type.shape[i] == in_types[0].shape[i])

        total_output_dim = sum([out_type.shape[self.axis].eval()
                                for out_type in out_types])
        type_check.expect(
            type_check.Variable(
                total_output_dim, 'sum(outputs[:][axis])'
            ) == in_types[0].shape[self.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:<commit_msg>Refactor type check of split axis<commit_after>import collections

import numpy
import six

from chainer import cuda
from chainer import function
from chainer.utils import type_check


_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 check_type_forward(self, in_types):
        type_check.expect(
            in_types.size() == 1,
            in_types[0].ndim >= self.axis
        )

        if isinstance(self.indices_or_sections, collections.Iterable):
            max_index = type_check.Variable(
                self.indices_or_sections[-1], 'max_index')
            type_check.expect(in_types[0].shape[self.axis] >= max_index)
        else:
            sections = type_check.Variable(
                self.indices_or_sections, 'sections')
            type_check.expect(in_types[0].shape[self.axis] % sections == 0)

    def check_type_backward(self, in_types, out_types):
        if isinstance(self.indices_or_sections, collections.Iterable):
            num_indices = type_check.Variable(
                len(self.indices_or_sections), 'len(indices)')
            output_size = num_indices + 1
        else:
            output_size = type_check.Variable(
                self.indices_or_sections, 'sections')

        type_check.expect(out_types.size() == output_size.eval())

        for out_type in out_types:
            type_check.expect(
                out_type.ndim == in_types[0].ndim,
                out_type.dtype == out_types[0].dtype
            )

        for out_type in out_types:
            for i in six.moves.range(in_types[0].ndim.eval()):
                if i == self.axis:
                    continue
                type_check.expect(out_type.shape[i] == in_types[0].shape[i])

        total_output_dim = sum(
            out_type.shape[self.axis] for out_type in out_types)
        type_check.expect(
            total_output_dim == in_types[0].shape[self.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:
{'fraction_non_alphanumeric': 0.10483637416739068, 'fraction_numerical': 0.006081668114682885, 'mean_word_length': 2.0945340501792113, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5501834', 'n_tokens_mistral': 2161, 'n_tokens_neox': 2092, 'n_words': 451}