From 6b453d5e4b120ed39c07d7d899bcd4939fcca163 Mon Sep 17 00:00:00 2001 From: Cookiee235 Date: Wed, 16 Apr 2025 21:27:00 +0800 Subject: [PATCH] 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 52d52e1936c2..3b18081b29fc 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