Searched refs:arg_max (Results 1 – 5 of 5) sorted by relevance
| /third_party/mindspore/tests/st/ops/graph_kernel/ |
| D | test_argmax.py | 27 self.arg_max = P.Argmax(axis=axis)
30 return self.arg_max(x)
|
| /third_party/mindspore/tests/vm_impl/ |
| D | vm_me.py | 445 arg_max = np.argmax(col, axis=1)
446 dmax[np.arange(arg_max.size), arg_max.flatten()] = dout.flatten()
453 def max_pool_grad_with_argmax(x, dout, arg_max, pool_h, pool_w, stride): argument
458 dmax[np.arange(arg_max.size), arg_max.flatten()] = dout.flatten()
|
| D | nn_ops_vm_impl.py | 136 arg_max = argmax.asnumpy()
137 dx = vm.max_pool_grad_with_argmax(x, dout, arg_max,
|
| /third_party/mindspore/mindspore/ops/_op_impl/cpu/ |
| D | __init__.py | 26 from .arg_max import _arg_max_cpu
|
| /third_party/mindspore/mindspore/ops/_op_impl/tbe/ |
| D | __init__.py | 274 from .arg_max import _arg_max_tbe
|