onnx--onnx/onnx/reference/ops/op_gather_elements.py

# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import numpy as np

from onnx.reference.op_run import OpRun


def gather_numpy_2(self: np.ndarray, index: np.ndarray) -> np.ndarray:
    res = []
    for a, b in zip(self, index, strict=True):
        res.append(a[b[0]])
    return np.array(res, dtype=self.dtype).reshape(index.shape)


def gather_numpy(self: np.ndarray, dim: int, index: np.ndarray) -> np.ndarray:
    idx_xsection_shape = index.shape[:dim] + index.shape[dim + 1 :]
    self_xsection_shape = self.shape[:dim] + self.shape[dim + 1 :]
    if idx_xsection_shape != self_xsection_shape:
        raise ValueError(
            f"Except for dimension {dim!r}, all dimensions of "
            f"index and self should be the same size."
        )
    data_swapped = np.swapaxes(self, 0, dim)
    index_swapped = np.swapaxes(index, 0, dim)

    try:
        gathered = np.choose(index_swapped, data_swapped, mode="wrap")
    except ValueError:
        if len(index_swapped.shape) == 2 and len(data_swapped.shape) == 2:
            return gather_numpy_2(self, index)
        raise  # pragma: no cover

    return np.swapaxes(gathered, 0, dim)


class GatherElements(OpRun):
    def _run(self, data, indices, axis=None):
        if indices.size == 0:
            return (np.empty((0,), dtype=data.dtype),)
        try:
            return (gather_numpy(data, axis, indices),)
        except TypeError:
            # distribution x86 requires int32.
            return (gather_numpy(data, axis, indices.astype(int)),)