43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
# 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 _argmax(data, axis=0, keepdims=True): # type: ignore
|
|
result = np.argmax(data, axis=axis)
|
|
if keepdims and len(result.shape) < len(data.shape):
|
|
result = np.expand_dims(result, axis)
|
|
return result.astype(np.int64)
|
|
|
|
|
|
def _argmax_use_numpy_select_last_index(data, axis=0, keepdims=True): # type: ignore
|
|
data = np.flip(data, axis)
|
|
result = np.argmax(data, axis=axis)
|
|
result = data.shape[axis] - result - 1
|
|
if keepdims:
|
|
result = np.expand_dims(result, axis)
|
|
return result.astype(np.int64)
|
|
|
|
|
|
class _ArgMax(OpRun):
|
|
def _run(self, data, axis=None, keepdims=None): # type: ignore
|
|
return (_argmax(data, axis=axis, keepdims=keepdims),)
|
|
|
|
|
|
class ArgMax_1(_ArgMax):
|
|
pass
|
|
|
|
|
|
class ArgMax_12(_ArgMax):
|
|
def _run(self, data, axis=None, keepdims=None, select_last_index=None): # type: ignore
|
|
if select_last_index == 0: # type: ignore
|
|
return _ArgMax._run(self, data, axis=axis, keepdims=keepdims)
|
|
return (
|
|
_argmax_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims),
|
|
)
|