<commit_before>import collections

import six

import chainer
from chainer.backends import cuda
from chainer.backends import intel64
from chainer import function_node
from chainer.utils import type_check


class SplitAxis(function_node.FunctionNode):

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

    def __init__(self, indices_or_sections, axis):
        if not isinstance(
                indices_or_sections,
                six.integer_types + (collections.Iterable,)):
            raise TypeError('indices_or_sections must be integer or 1-D array')
        if (chainer.is_debug() and
                isinstance(indices_or_sections, collections.Iterable)):
            for p, n in six.moves.zip(
                    indices_or_sections, indices_or_sections[1:]):
                if p > n:
                    raise ValueError('indices_or_sections must be sorted')
        self.indices_or_sections = indices_or_sections
        self.axis = axis

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() == 1)
        type_check.expect(in_types[0].ndim > self.axis)

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

    def forward(self, inputs):
        # Currently iDeep only supports 4 dims
        if (intel64.should_use_ideep('>=auto')
                and intel64.inputs_all_ready(inputs, (4,))
                and self._ideep_is_indices_or_sections_supported(
                    self.indices_or_sections)):
            return self._forward_ideep(inputs)

        x, = inputs
        if isinstance(self.indices_or_sections, collections.Iterable):
            cdimx = x.shape[self.axis]
            ind = list(self.indices_or_sections)
            ind.append(cdimx)
        self._xp = cuda.get_array_module(x)
        ret = tuple(self._xp.split(x, self.indices_or_sections, self.axis))
        self._shapes = [r.shape for r in ret]
        return ret

    def _ideep_is_indices_or_sections_supported(self, indices_or_sections):
        # Returns True if the value of ``indices_or_sections`` is supported
        # by iDeep.
        ios = indices_or_sections
        if isinstance(ios, collections.Iterable):
            if len(ios) == 0:
                return False  # Empty sequence
            if ios[0] == 0:
                return False  # Sequence starting with 0
            if any([ios[i-1] == ios[i] for i in six.moves.range(1, len(ios))]):
                return False  # Sequence with duplicate index
        else:
            if ios == 1:
                return False  # 1
        return True

    def _forward_ideep(self, inputs):
        x, = inputs
        offsets = intel64.ideep.intVector()
        # TODO(iDeep)
        # bypass python3 issue when transfer array to std::vector<>
        # https://github.com/SimpleITK/SimpleITK/issues/106
        ios = self.indices_or_sections
        if isinstance(ios, collections.Iterable):
            for i in ios:
                offsets.push_back(i)
        else:
            d = x.shape[self.axis]<commit_msg>F.split_axis: use for loop instead of any()<commit_after>import collections

import six

import chainer
from chainer.backends import cuda
from chainer.backends import intel64
from chainer import function_node
from chainer.utils import type_check


class SplitAxis(function_node.FunctionNode):

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

    def __init__(self, indices_or_sections, axis):
        if not isinstance(
                indices_or_sections,
                six.integer_types + (collections.Iterable,)):
            raise TypeError('indices_or_sections must be integer or 1-D array')
        if (chainer.is_debug() and
                isinstance(indices_or_sections, collections.Iterable)):
            for p, n in six.moves.zip(
                    indices_or_sections, indices_or_sections[1:]):
                if p > n:
                    raise ValueError('indices_or_sections must be sorted')
        self.indices_or_sections = indices_or_sections
        self.axis = axis

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() == 1)
        type_check.expect(in_types[0].ndim > self.axis)

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

    def forward(self, inputs):
        # Currently iDeep only supports 4 dims
        if (intel64.should_use_ideep('>=auto')
                and intel64.inputs_all_ready(inputs, (4,))
                and self._ideep_is_indices_or_sections_supported(
                    self.indices_or_sections)):
            return self._forward_ideep(inputs)

        x, = inputs
        if isinstance(self.indices_or_sections, collections.Iterable):
            cdimx = x.shape[self.axis]
            ind = list(self.indices_or_sections)
            ind.append(cdimx)
        self._xp = cuda.get_array_module(x)
        ret = tuple(self._xp.split(x, self.indices_or_sections, self.axis))
        self._shapes = [r.shape for r in ret]
        return ret

    def _ideep_is_indices_or_sections_supported(self, indices_or_sections):
        # Returns True if the value of ``indices_or_sections`` is supported
        # by iDeep.
        ios = indices_or_sections
        if isinstance(ios, collections.Iterable):
            if len(ios) == 0:
                return False  # Empty sequence
            if ios[0] == 0:
                return False  # Sequence starting with 0
            for i in six.moves.range(1, len(ios)):
                if ios[i-1] == ios[i]:
                    return False  # Sequence with duplicate index
        else:
            if ios == 1:
                return False  # 1
        return True

    def _forward_ideep(self, inputs):
        x, = inputs
        offsets = intel64.ideep.intVector()
        # TODO(iDeep)
        # bypass python3 issue when transfer array to std::vector<>
        # https://github.com/SimpleITK/SimpleITK/issues/106
        ios = self.indices_or_sections
        if isinstance(ios, collections.Iterable):
            for i in ios:
                offsets.push_back(i)
        else:
            d = x.shape[self.axis]
{'fraction_non_alphanumeric': 0.08802616609783845, 'fraction_numerical': 0.00881683731513083, 'mean_word_length': 2.145348837209302, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 15, 'https://': 2, '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': '4136075', 'n_tokens_mistral': 2074, 'n_tokens_neox': 1971, 'n_words': 471}