Spaces:
Sleeping
Sleeping
File size: 1,270 Bytes
dc2106c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
# 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),
)
|