def argmax(input, axis=None, name=None, dimension=None):
if dimension is not None:
if axis is not None:
raise ValueError("cannot specify both 'axis' and 'dimension'.")
axis = dimension
elif axis is None: axis = 0
return ops.Argmax(input, axis=axis, name=name)
def argmax(x, axis=None, keepdims=False):
"""Compute the indices of maximum elements along the given axis.
Parameters
----------
x : Tensor
The input tensor.
axis : int
The axis to compute. Default is ``None`` (Along all axes).
keep_dims : boolean
Whether to keep dims after computing.
Returns
-------
Tensor
The indices.
"""
if axis is None: axis = -1
return ops.Argmax(x, axis=axis, keep_dims=keepdims)
def Setup(self, bottom):
super(ArgMaxLayer, self).Setup(bottom)
input = bottom[0] if isinstance(bottom, list) else bottom
return ops.Argmax(input, **self._param)
