(tvm) branch main updated: Fix incorrect docstring in topi softmax (#17844)

[mshr]

This is an automated email from the ASF dual-hosted git repository.

mshr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 40a16db73e Fix incorrect docstring in topi softmax  (#17844)
40a16db73e is described below

commit 40a16db73ef0d221dbebd17c8ec49039ff898e33
Author: Cookiee235 <qingc...@tju.edu.cn>
AuthorDate: Fri Apr 18 16:14:33 2025 +0800

    Fix incorrect docstring in topi softmax  (#17844)
    
    Update softmax.py
    
    Fix the inconsistent docstring with parameter
---
 python/tvm/topi/nn/softmax.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/python/tvm/topi/nn/softmax.py b/python/tvm/topi/nn/softmax.py
index 52d52e1936..3b18081b29 100644
--- a/python/tvm/topi/nn/softmax.py
+++ b/python/tvm/topi/nn/softmax.py
@@ -27,7 +27,7 @@ def softmax(x, axis=-1):
 
     Parameters
     ----------
-    data : tvm.te.Tensor
+    x : tvm.te.Tensor
         can be any dimension
 
     axis : int
@@ -48,7 +48,7 @@ def fast_softmax(x, axis=-1):
 
     Parameters
     ----------
-    data : tvm.te.Tensor
+    x : tvm.te.Tensor
         can be any dimension
 
     axis : int
@@ -128,9 +128,12 @@ def log_softmax(x, axis=-1):
 
     Parameters
     ----------
-    data : tvm.te.Tensor
+    x : tvm.te.Tensor
         N-D input data
 
+    axis : int
+        channel axis
+
     Returns
     -------
     output : tvm.te.Tensor