Skip to content

Commit

Permalink
correct arg passing
Browse files Browse the repository at this point in the history
  • Loading branch information
AndrewZhaoLuo authored and Andrew Zhao Luo committed Aug 27, 2021
1 parent 20cdd36 commit 49b6322
Showing 1 changed file with 12 additions and 4 deletions.
16 changes: 12 additions & 4 deletions python/tvm/topi/reduction.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ def min(data, axis=None, keepdims=False):
return cpp.min(data, axis, keepdims)


def argmax(data, axis=None, keepdims=False, exclude=False, select_last_index=False):
def argmax(data, axis=None, keepdims=False, select_last_index=False):
"""Returns the indices of the maximum values along an axis.
Parameters
Expand All @@ -185,14 +185,18 @@ def argmax(data, axis=None, keepdims=False, exclude=False, select_last_index=Fal
with size one.
With this option, the result will broadcast correctly against the input array.
select_last_index: bool
Whether to select the last index if the maximum element appears multiple times, else
select the first index.
Returns
-------
ret : tvm.te.Tensor
"""
return cpp.argmax(data, axis, keepdims, exclude=exclude, select_last_index=select_last_index)
return cpp.argmax(data, axis, keepdims, select_last_index)


def argmin(data, axis=None, keepdims=False, exclude=False, select_last_index=False):
def argmin(data, axis=None, keepdims=False, select_last_index=False):
"""Returns the indices of the minimum values along an axis.
Parameters
Expand All @@ -210,11 +214,15 @@ def argmin(data, axis=None, keepdims=False, exclude=False, select_last_index=Fal
with size one.
With this option, the result will broadcast correctly against the input array.
select_last_index: bool
Whether to select the last index if the minimum element appears multiple times, else
select the first index.
Returns
-------
ret : tvm.te.Tensor
"""
return cpp.argmin(data, axis, keepdims, exclude, select_last_index)
return cpp.argmin(data, axis, keepdims, select_last_index)


def prod(data, axis=None, keepdims=False):
Expand Down

0 comments on commit 49b6322

Please sign in to comment.