From 11e2ed1a0b8a83e225eeaf8958995a4c964dd82c Mon Sep 17 00:00:00 2001 From: Cookiee235 Date: Thu, 17 Apr 2025 15:11:38 +0800 Subject: [PATCH 1/2] Update batch_to_space_nd.py --- python/tvm/topi/nn/batch_to_space_nd.py | 2 +- 1 file changed, 1 insertion(+), 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 c61a90a7777b..99e6312f402e 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. From 27cfe8e66b509470eabd640fffb3862b6311a019 Mon Sep 17 00:00:00 2001 From: Cookiee235 Date: Thu, 17 Apr 2025 15:41:48 +0800 Subject: [PATCH 2/2] Update bitserial_util.py --- python/tvm/topi/nn/bitserial_util.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/tvm/topi/nn/bitserial_util.py b/python/tvm/topi/nn/bitserial_util.py index 3a55422493d4..8b9de94b3e43 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)