<commit_before>import warnings

import chainer
import numpy as np
from onnx.mapping import NP_TYPE_TO_TENSOR_TYPE

from onnx_chainer.functions.opset_version import support
from onnx_chainer import onnx_helper


TENSOR_TYPE_TO_NAME = {
    0: 'UNDEFINED',
    1: 'FLOAT',
    2: 'UINT8',
    3: 'INT8',
    4: 'UINT16',
    5: 'INT16',
    6: 'INT32',
    7: 'INT64',
    8: 'STRING',
    9: 'BOOL',
    10: 'FLOAT16',
    11: 'DOUBLE',
    12: 'UINT32',
    13: 'UINT64',
    14: 'COMPLEX64',
    15: 'COMPLEX128',
}


@support((1, 6))
def convert_Cast(func, opset_version, input_names, output_names, context):
    typ = func.type if isinstance(func.type, np.dtype) else np.dtype(func.type)
    if opset_version == 1:
        return onnx_helper.make_node(
            'Cast', input_names, output_names,
            to=TENSOR_TYPE_TO_NAME[NP_TYPE_TO_TENSOR_TYPE[typ]]
        ),
    elif opset_version == 6:
        return onnx_helper.make_node(
            'Cast', input_names, output_names,
            to=NP_TYPE_TO_TENSOR_TYPE[typ]
        ),


@support((1, 4))
def convert_Concat(func, opset_version, input_names, output_names, context):
    if opset_version == 1:
        return onnx_helper.make_node(
            'Concat', input_names, output_names,
            axis=func.axis
        ),
    elif opset_version == 4:
        return onnx_helper.make_node(
            'Concat', input_names, output_names,
            axis=func.axis
        ),


def convert_Copy(func, opset_version, input_names, output_names, context):
    return onnx_helper.make_node(
        'Identity', input_names, output_names
    ),


def convert_Depth2Space(
        func, opset_version, input_names, output_names, context):
    return onnx_helper.make_node(
        'DepthToSpace', input_names, output_names,
        blocksize=func.r
    ),


def get_slice_node(
        gb, opset_version, context, input_names, axes, starts, ends):
    if opset_version < 10:
        return gb.op(
            'Slice', input_names, axes=axes, starts=starts, ends=ends)
    else:
        for param in [('starts', starts), ('ends', ends), ('axes', axes)]:
            param_name = context.add_const(
                np.asarray(list(param[1]), dtype=np.int64), param[0])
            input_names.append(param_name)
        return gb.op('Slice', input_names)


def _to_ndarray(x, dtype=np.int64):
    if isinstance(x, list):
        return np.array(x, dtype=dtype)
    else:
        return chainer.cuda.to_cpu(x).astype(dtype)


@support((1, 10))
def convert_GetItem(func, opset_version, input_names, output_names, context):
    x = func.inputs[0]
    axes, starts, ends = [], [], []
    squeeze_idxs, unsqueeze_idxs = [], []
    skipped = 0  # when set ellipsis, need to skip index rolling

    prev_gathered_axis = -1
    gather_axis = -1
    gather_idx = None  # when GatherND, set first array for broadcasting
    gather_nd_idx = None
    is_used_slice_whole = False  # GatherND does not support axis, need to care

    for i, idx in enumerate(func.slices):
        # axis means the index of input x, adjust None and Ellipsis counts
        axis = i - len(unsqueeze_idxs) + skipped
        if isinstance(idx, slice):
            if idx.step is not None and idx.step != 1:
                raise ValueError(
                    'GetItem with {}step slicing is not supported in ONNX '
                    'Slice operator'.format(idx.step))
            if idx.start is None and idx.stop is None:
                is_used_slice_whole = True
                continue
            axes.append(axis)
            starts.append(0 if idx.start is None else idx.start)
            ends.append(x.shape[axis] if idx.stop is None else idx.stop)
        elif isinstance(idx, int):
            axes.append(axis)
            starts.append(idx)
            ends.append(idx+1)
            squeeze_idxs.append(axis)
        elif isinstance(idx, np.ndarray) and idx.ndim == 0:
            scalar_idx = idx.item()
            axes.append(axis)
            starts.append(scalar_idx)
            ends.append(scalar_idx+1)
            squeeze_idxs.append(axis)
        elif idx is None:
            unsqueeze_idxs.append(i - len(squeeze_idxs) + skipped)
        elif idx is Ellipsis:
            # calculate rest slice number except None, GetItem does not allow
            # multiple Ellipsis, so ignore latter Ellipsis count
            rest_slice_len = len(
                [idx_ for idx_ in func.slices[i+1:] if idx_ is not None])
            assert skipped == 0
            skipped = len(x.shape) - axis - rest_slice_len - 1
        elif isinstance(idx, (list,) + chainer.get_array_types()):
            if prev_gathered_axis >= 0:
                if (i-1) != prev_gathered_axis:
                    raise ValueError(
                        'ONNX-Chainer does not support split multiple advanced'
                        ' indexing')
                if is_used_slice_whole:
                    raise ValueError(
                        'ONNX-Chainer does not support whole indexing(`[:]`)'
                        'in front of multiple advanced indexing')
                if unsqueeze_idxs:
                    raise ValueError(
                        'ONNX-Chainer does not support new axis in front of '
                        'multiple advanced indexing')
                # multiple advanced index, convert to GatherND
                idx_array = _to_ndarray(idx)
                base_idx = gather_idx if gather_nd_idx is None else\
                    gather_nd_idx
                gather_nd_idx = np.vstack((base_idx, idx_array))
                prev_gathered_axis = i
            else:
                # convert to Gather, if next index is also list, change to
                # GatherND
                gather_axis = axis - len(squeeze_idxs) + len(unsqueeze_idxs)
                gather_idx = _to_ndarray(idx)<commit_msg>fix error message<commit_after>import warnings

import chainer
import numpy as np
from onnx.mapping import NP_TYPE_TO_TENSOR_TYPE

from onnx_chainer.functions.opset_version import support
from onnx_chainer import onnx_helper


TENSOR_TYPE_TO_NAME = {
    0: 'UNDEFINED',
    1: 'FLOAT',
    2: 'UINT8',
    3: 'INT8',
    4: 'UINT16',
    5: 'INT16',
    6: 'INT32',
    7: 'INT64',
    8: 'STRING',
    9: 'BOOL',
    10: 'FLOAT16',
    11: 'DOUBLE',
    12: 'UINT32',
    13: 'UINT64',
    14: 'COMPLEX64',
    15: 'COMPLEX128',
}


@support((1, 6))
def convert_Cast(func, opset_version, input_names, output_names, context):
    typ = func.type if isinstance(func.type, np.dtype) else np.dtype(func.type)
    if opset_version == 1:
        return onnx_helper.make_node(
            'Cast', input_names, output_names,
            to=TENSOR_TYPE_TO_NAME[NP_TYPE_TO_TENSOR_TYPE[typ]]
        ),
    elif opset_version == 6:
        return onnx_helper.make_node(
            'Cast', input_names, output_names,
            to=NP_TYPE_TO_TENSOR_TYPE[typ]
        ),


@support((1, 4))
def convert_Concat(func, opset_version, input_names, output_names, context):
    if opset_version == 1:
        return onnx_helper.make_node(
            'Concat', input_names, output_names,
            axis=func.axis
        ),
    elif opset_version == 4:
        return onnx_helper.make_node(
            'Concat', input_names, output_names,
            axis=func.axis
        ),


def convert_Copy(func, opset_version, input_names, output_names, context):
    return onnx_helper.make_node(
        'Identity', input_names, output_names
    ),


def convert_Depth2Space(
        func, opset_version, input_names, output_names, context):
    return onnx_helper.make_node(
        'DepthToSpace', input_names, output_names,
        blocksize=func.r
    ),


def get_slice_node(
        gb, opset_version, context, input_names, axes, starts, ends):
    if opset_version < 10:
        return gb.op(
            'Slice', input_names, axes=axes, starts=starts, ends=ends)
    else:
        for param in [('starts', starts), ('ends', ends), ('axes', axes)]:
            param_name = context.add_const(
                np.asarray(list(param[1]), dtype=np.int64), param[0])
            input_names.append(param_name)
        return gb.op('Slice', input_names)


def _to_ndarray(x, dtype=np.int64):
    if isinstance(x, list):
        return np.array(x, dtype=dtype)
    else:
        return chainer.cuda.to_cpu(x).astype(dtype)


@support((1, 10))
def convert_GetItem(func, opset_version, input_names, output_names, context):
    x = func.inputs[0]
    axes, starts, ends = [], [], []
    squeeze_idxs, unsqueeze_idxs = [], []
    skipped = 0  # when set ellipsis, need to skip index rolling

    prev_gathered_axis = -1
    gather_axis = -1
    gather_idx = None  # when GatherND, set first array for broadcasting
    gather_nd_idx = None
    is_used_slice_whole = False  # GatherND does not support axis, need to care

    for i, idx in enumerate(func.slices):
        # axis means the index of input x, adjust None and Ellipsis counts
        axis = i - len(unsqueeze_idxs) + skipped
        if isinstance(idx, slice):
            if idx.step is not None and idx.step != 1:
                raise ValueError(
                    'GetItem with {}step slicing is not supported in ONNX '
                    'Slice operator'.format(idx.step))
            if idx.start is None and idx.stop is None:
                is_used_slice_whole = True
                continue
            axes.append(axis)
            starts.append(0 if idx.start is None else idx.start)
            ends.append(x.shape[axis] if idx.stop is None else idx.stop)
        elif isinstance(idx, int):
            axes.append(axis)
            starts.append(idx)
            ends.append(idx+1)
            squeeze_idxs.append(axis)
        elif isinstance(idx, np.ndarray) and idx.ndim == 0:
            scalar_idx = idx.item()
            axes.append(axis)
            starts.append(scalar_idx)
            ends.append(scalar_idx+1)
            squeeze_idxs.append(axis)
        elif idx is None:
            unsqueeze_idxs.append(i - len(squeeze_idxs) + skipped)
        elif idx is Ellipsis:
            # calculate rest slice number except None, GetItem does not allow
            # multiple Ellipsis, so ignore latter Ellipsis count
            rest_slice_len = len(
                [idx_ for idx_ in func.slices[i+1:] if idx_ is not None])
            assert skipped == 0
            skipped = len(x.shape) - axis - rest_slice_len - 1
        elif isinstance(idx, (list,) + chainer.get_array_types()):
            if prev_gathered_axis >= 0:
                if (i - 1) != prev_gathered_axis:
                    raise ValueError(
                        'ONNX-Chainer does not support non-consecutive'
                        'multiple advanced indexing')
                if is_used_slice_whole:
                    raise ValueError(
                        'ONNX-Chainer does not support whole indexing(`[:]`)'
                        'in front of multiple advanced indexing')
                if unsqueeze_idxs:
                    raise ValueError(
                        'ONNX-Chainer does not support new axis in front of '
                        'multiple advanced indexing')
                # multiple advanced index, convert to GatherND
                idx_array = _to_ndarray(idx)
                base_idx = gather_idx if gather_nd_idx is None else\
                    gather_nd_idx
                gather_nd_idx = np.vstack((base_idx, idx_array))
                prev_gathered_axis = i
            else:
                # convert to Gather, if next index is also list, change to
                # GatherND
                gather_axis = axis - len(squeeze_idxs) + len(unsqueeze_idxs)
                gather_idx = _to_ndarray(idx)
{'fraction_non_alphanumeric': 0.09397057572922868, 'fraction_numerical': 0.013096351730589336, 'mean_word_length': 2.0505836575875485, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '102413', 'n_tokens_mistral': 3746, 'n_tokens_neox': 3584, 'n_words': 934}