<commit_before># distutils: language = c++
import sys

import numpy

import cupy
from cupy.core._kernel import ElementwiseKernel
from cupy.core._ufuncs import elementwise_copy

from libcpp cimport vector

from cupy.core cimport core
from cupy.core cimport _routines_math as _math
from cupy.core cimport _routines_manipulation as _manipulation
from cupy.core.core cimport ndarray
from cupy.core cimport internal


# ndarray members


cdef ndarray _ndarray_getitem(ndarray self, slices):
    # supports basic indexing (by slices, ints or Ellipsis) and
    # some parts of advanced indexing by integer or boolean arrays.
    # TODO(beam2d): Support the advanced indexing of NumPy.
    cdef Py_ssize_t mask_i
    cdef list slice_list, adv_mask, adv_slices
    cdef bint advanced, mask_exists

    slice_list, advanced, mask_exists = _prepare_slice_list(
        slices, self._shape.size())

    if mask_exists:
        mask_i = _get_mask_index(slice_list)
        return _getitem_mask_single(self, slice_list[mask_i], mask_i)
    if advanced:
        a, adv_slices, adv_mask = _prepare_advanced_indexing(
            self, slice_list)
        if sum(adv_mask) == 1:
            axis = adv_mask.index(True)
            return a.take(adv_slices[axis], axis)
        return _getitem_multiple(a, adv_slices)

    return _simple_getitem(self, slice_list)


cdef _ndarray_setitem(ndarray self, slices, value):
    _scatter_op(self, slices, value, 'update')


cdef tuple _ndarray_nonzero(ndarray self):
    cdef int ndim
    cdef ndarray dst = _ndarray_argwhere(self)
    ndim = len(dst)
    if ndim > 1:
        return tuple([dst[i] for i in range(ndim)])
    elif ndim == 1:
        return (dst,)


cdef ndarray _ndarray_argwhere(ndarray self):
    cdef Py_ssize_t count_nonzero
    cdef int ndim
    dtype = numpy.int64
    if self.size == 0:
        count_nonzero = 0
    else:
        r = self.ravel()
        nonzero = cupy.core.not_equal(r, 0, ndarray(r.shape, dtype))
        del r
        scan_index = _math.scan(nonzero, op=_math.scan_op.SCAN_SUM)
        count_nonzero = int(scan_index[-1])  # synchronize!
    ndim = max(<int>self._shape.size(), 1)
    if count_nonzero == 0:
        return ndarray((0, ndim), dtype=dtype)

    if ndim <= 1:
        dst = ndarray((count_nonzero, 1), dtype=dtype)
        _nonzero_kernel_1d(nonzero, scan_index, dst)
        return dst
    else:
        nonzero.shape = self.shape
        scan_index.shape = self.shape
        dst = ndarray((count_nonzero, ndim), dtype=dtype)
        _nonzero_kernel(nonzero, scan_index, dst)
        return dst


cdef _ndarray_scatter_add(ndarray self, slices, value):
    _scatter_op(self, slices, value, 'add')


cdef _ndarray_scatter_max(ndarray self, slices, value):
    _scatter_op(self, slices, value, 'max')


cdef _ndarray_scatter_min(ndarray self, slices, value):
    _scatter_op(self, slices, value, 'min')


cdef ndarray _ndarray_take(ndarray self, indices, axis, out):
    if axis is None:<commit_msg>drop ndim <= 1 case<commit_after># distutils: language = c++
import sys

import numpy

import cupy
from cupy.core._kernel import ElementwiseKernel
from cupy.core._ufuncs import elementwise_copy

from libcpp cimport vector

from cupy.core cimport core
from cupy.core cimport _routines_math as _math
from cupy.core cimport _routines_manipulation as _manipulation
from cupy.core.core cimport ndarray
from cupy.core cimport internal


# ndarray members


cdef ndarray _ndarray_getitem(ndarray self, slices):
    # supports basic indexing (by slices, ints or Ellipsis) and
    # some parts of advanced indexing by integer or boolean arrays.
    # TODO(beam2d): Support the advanced indexing of NumPy.
    cdef Py_ssize_t mask_i
    cdef list slice_list, adv_mask, adv_slices
    cdef bint advanced, mask_exists

    slice_list, advanced, mask_exists = _prepare_slice_list(
        slices, self._shape.size())

    if mask_exists:
        mask_i = _get_mask_index(slice_list)
        return _getitem_mask_single(self, slice_list[mask_i], mask_i)
    if advanced:
        a, adv_slices, adv_mask = _prepare_advanced_indexing(
            self, slice_list)
        if sum(adv_mask) == 1:
            axis = adv_mask.index(True)
            return a.take(adv_slices[axis], axis)
        return _getitem_multiple(a, adv_slices)

    return _simple_getitem(self, slice_list)


cdef _ndarray_setitem(ndarray self, slices, value):
    _scatter_op(self, slices, value, 'update')


cdef tuple _ndarray_nonzero(ndarray self):
    cdef int ndim
    cdef ndarray dst = _ndarray_argwhere(self)
    ndim = len(dst)
    if ndim > 1:
        return tuple([dst[i] for i in range(ndim)])
    elif ndim == 1:
        return (dst,)


cdef ndarray _ndarray_argwhere(ndarray self):
    cdef Py_ssize_t count_nonzero
    cdef int ndim
    dtype = numpy.int64
    if self.size == 0:
        count_nonzero = 0
    else:
        r = self.ravel()
        nonzero = cupy.core.not_equal(r, 0, ndarray(r.shape, dtype))
        del r
        scan_index = _math.scan(nonzero, op=_math.scan_op.SCAN_SUM)
        count_nonzero = int(scan_index[-1])  # synchronize!
    ndim = max(<int>self._shape.size(), 1)
    if count_nonzero == 0:
        return ndarray((0, ndim), dtype=dtype)

    nonzero.shape = self.shape
    scan_index.shape = self.shape
    dst = ndarray((count_nonzero, ndim), dtype=dtype)
    _nonzero_kernel(nonzero, scan_index, dst)
    return dst


cdef _ndarray_scatter_add(ndarray self, slices, value):
    _scatter_op(self, slices, value, 'add')


cdef _ndarray_scatter_max(ndarray self, slices, value):
    _scatter_op(self, slices, value, 'max')


cdef _ndarray_scatter_min(ndarray self, slices, value):
    _scatter_op(self, slices, value, 'min')


cdef ndarray _ndarray_take(ndarray self, indices, axis, out):
    if axis is None:
{'fraction_non_alphanumeric': 0.08766177739430543, 'fraction_numerical': 0.005176876617773943, 'mean_word_length': 3.377643504531722, 'pattern_counts': {'":': 0, '<': 7, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2247327', 'n_tokens_mistral': 2061, 'n_tokens_neox': 1984, 'n_words': 529}