(tvm) branch main updated: Fix docstring in batch_to_space_nd and bitpack (#17848)

[syfeng]

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

syfeng 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 982b46c3b3 Fix docstring in batch_to_space_nd and bitpack (#17848)
982b46c3b3 is described below

commit 982b46c3b38666ff339a5ab7e79cb5947bba8342
Author: Cookiee235 <qingc...@tju.edu.cn>
AuthorDate: Fri Apr 18 09:58:29 2025 +0800

    Fix docstring in batch_to_space_nd and bitpack (#17848)
    
    * Update batch_to_space_nd.py
    
    * Update bitserial_util.py
---
 python/tvm/topi/nn/batch_to_space_nd.py | 2 +-
 python/tvm/topi/nn/bitserial_util.py    | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/python/tvm/topi/nn/batch_to_space_nd.py 
b/python/tvm/topi/nn/batch_to_space_nd.py
index c61a90a777..99e6312f40 100644
--- a/python/tvm/topi/nn/batch_to_space_nd.py
+++ b/python/tvm/topi/nn/batch_to_space_nd.py
@@ -29,7 +29,7 @@ def batch_to_space_nd(data, block_shape, crop_begin_list, 
crop_end_list):
         N-D Tensor with shape [batch, spatial_shape, remaining_shapes],
         where spatial_shape has M dimensions.
 
-    block_size : list of ints
+    block_shape : list of ints
         list of size [M] where M is number of spatial dims, specifies block
         size for each spatial dimension.
 
diff --git a/python/tvm/topi/nn/bitserial_util.py 
b/python/tvm/topi/nn/bitserial_util.py
index 3a55422493..8b9de94b3e 100644
--- a/python/tvm/topi/nn/bitserial_util.py
+++ b/python/tvm/topi/nn/bitserial_util.py
@@ -28,10 +28,18 @@ def bitpack(data, bits, pack_axis, bit_axis, pack_type, 
name="QuantizeInput"):
 
     Parameters
     ----------
+    data : tvm.te.Tensor
+        The input tvm tensor
+    bits : int
+        Number of bits to use for packing
     pack_axis : int
        index of the axis to pack in data
     bit_axis : int
        index of axis to place bit axis in resulting packed data
+    pack_type : str
+        Data type for packing, must be one of: ['uint8', 'uint16', 'uint32', 
'uint64']
+    name : Optional[str] = "QuantizeInput"
+        Name for the operation
     """
     ishape = data.shape
     n = len(ishape)