Spaces:
Sleeping
Sleeping
| # Copyright (c) ONNX Project Contributors | |
| # SPDX-License-Identifier: Apache-2.0 | |
| 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), | |
| ) | |