v2.1.17: fix a bug in sparse add, more asserts

CHANGELOG.md

 # Changelog
 
+## [2.1.17] - 2021-11-29
+### Fixed
+- Fix a bug in sparse add.
+### Added
+- Add more wrong usage check
+- Add insert_exist_keys for hash table
+
 ## [2.1.16] - 2021-11-28
 ### Fixed
 - Fix strange compile problem in windows

docs/USAGE.md

@@ -16,6 +16,47 @@
 
 # Usage
 
+## Short API description
+
+```Python
+import spconv.pytorch as spconv
+from spconv.pytorch import functional as Fsp
+from torch import nn
+from spconv.pytorch.utils import PointToVoxel
+from spconv.pytorch.hash import HashTable
+```
+
+| Layer APIs                         | Common Usage             |            Dense Version    |Note    |
+|----------------------------------- |:------------------------:|----------------------------:|----------------------------:| 
+| ```spconv.SparseConv3d```          | Downsample               | ```nn.Conv3d```             | Use ```indice_key``` to save data for inverse |
+| ```spconv.SubMConv3d```            | Convolution              | N/A                         | Use ```indice_key``` to save data for reuse |
+| ```spconv.SparseInverseConv3d```   | Upsample                 |  N/A                        | Use pre-saved ```indice_key``` to upsample |
+| ```spconv.SparseConvTranspose3d``` | Upsample (don't use this)|  ```nn.ConvTranspose3d```   | VERY SLOW and CAN'T RECOVER ORIGIN POINT CLOUD |
+| ```spconv.SparseMaxPool3d```       | Downsample               |  ```nn.MaxPool3d```         | Use ```indice_key``` to save data for inverse |
+| ```spconv.SparseSequential```       | Container               |  ```nn.Sequential```         | support layers above and ```nn.ReLU, nn.BatchNorm, ...```|
+
+
+| Functional APIs                    | Usage                    |
+|----------------------------------- |:------------------------:|
+| ```Fsp.sparse_add```               | Add sparse tensors with same shape and different indices    |
+
+| Input APIs                         | Usage                    |
+|----------------------------------- |:------------------------:|
+| ```PointToVoxel```                 | point cloud to voxels    |
+
+| Misc APIs                         | Usage                    |
+|----------------------------------- |:------------------------:|
+| ```HashTable```                   | hash table, one-slot     |
+
+| Layer APIs                         | [torchsparse](https://github.com/mit-han-lab/torchsparse)             |    [MinkowskiEngine](https://github.com/NVIDIA/MinkowskiEngine)             |   
+|----------------------------------- |:------------------------:|:------------------------:|
+| ```spconv.SparseConv3d```          | ```Conv3d(stride!=1, transpose=False)```               |```MinkowskiConvolution(stride!=1)```| 
+| ```spconv.SubMConv3d```            | ```Conv3d(stride=1, transpose=False)```              | ```MinkowskiConvolution(stride=1)```| 
+| ```spconv.SparseInverseConv3d```   | ```Conv3d(stride!=1, transpose=True)```                 |```MinkowskiConvolutionTranspose```| 
+| ```spconv.SparseConvTranspose3d``` | N/A                |```MinkowskiConvolutionTranspose```| 
+| ```spconv.SparseMaxPool3d```       | N/A               | ```MinkowskiMaxPooling```| 
+
+
 ## Concept
 
 * Sparse Conv Tensor: like hybird [torch.sparse_coo_tensor](https://pytorch.org/docs/stable/sparse.html#sparse-coo-docs) but only have two difference: 1. SparseConvTensor only have one dense dim, 2. indice of SparseConvTensor is transposed. see torch doc for more details.
@@ -102,6 +143,29 @@ class ExampleNet(nn.Module):
         return self.net(x)
 ```
 
+### Sparse Add
+
+In sematic segmentation network, we may use conv1x3, 3x1 and 3x3 in a block, but it's impossible to sum result from these layers because regular add requires same indices.
+
+spconv >= 2.1.17 provide a operation to add sparse tensors with different indices (shape must same), but with limits:
+
+```Python
+from spconv.pytorch import functional as Fsp
+res_1x3 = conv1x3(x)
+res_3x1 = conv3x1(x)
+# WRONG
+# because we can't "inverse" this operation
+wrong_usage_cant_inverse = Fsp.sparse_add(res_1x3, res_3x1)
+
+# CORRECT
+# res_3x3 already contains all indices of res_1x3 and res_3x1, 
+# so output spatial structure isn't changed, we can "inverse" back.
+res_3x3 = conv3x3(x)
+correct = Fsp.sparse_add(res_1x3, res_3x1, res_3x3)
+```
+
+If you use a network without ```SparseInverseConv```, limits above aren't exists, the only drawback of ```sparse_add``` is that it run slower than simple aligned add.
+
 ### Fast Mixed Percision Training
 
 see example/mnist_sparse. we support ```torch.cuda.amp```.

example/simple_hash.py

@@ -52,5 +52,12 @@ def main():
         cnt_item = cnt.item()
         print(cnt, ks[:cnt_item], vs[:cnt_item])
 
+        print("----------Insert Exist Keys----------")
+        is_empty = table.insert_exist_keys(keys, values)
+        ks, vs, cnt = table.items()
+        cnt_item = cnt.item()
+        print(cnt, ks[:cnt_item], vs[:cnt_item])
+
+
 if __name__ == "__main__":
     main()
\ No newline at end of file

spconv/core_cc/csrc/hash/core.pyi

@@ -74,3 +74,14 @@ class HashTable:
             stream: 
         """
         ...
+    def insert_exist_keys(self, keys: Tensor, values: Tensor, is_empty: Tensor, stream: int) -> None: 
+        """
+        insert v of given k if k exists. won't insert any new key.
+                
+        Args:
+            keys: 
+            values: 
+            is_empty: 
+            stream: 
+        """
+        ...

spconv/csrc/hash/core.py

@@ -54,6 +54,33 @@ def _dispatch(code: pccm.FunctionCode, dts: List[dtypes.DType], var: str):
         TV_THROW_RT_ERR("unknown dtype {var}, available: {dts}")
         """)
 
+class HashTableKernel(pccm.Class):
+    def __init__(self):
+        super().__init__()
+        self.add_dependency(TensorView, TensorViewHashKernel, TensorViewKernel)
+
+    @pccm.cuda.cuda_global_function
+    def insert_exist_keys_kernel(self):
+        code = pccm.FunctionCode()
+        code.targ("THashTableSplit")
+        code.arg("table", "THashTableSplit")
+        code.arg("key_ptr", "const typename THashTableSplit::key_type *__restrict__")
+        code.arg("value_ptr", "const typename THashTableSplit::mapped_type *__restrict__")
+        code.arg("is_empty_ptr", "uint8_t*")
+        code.arg("size", "size_t")
+
+        code.raw(f"""
+        auto value_data = table.value_ptr();
+        for (size_t i : tv::KernelLoopX<size_t>(size)){{
+            auto key = key_ptr[i];
+            auto offset = table.lookup_offset(key);
+            is_empty_ptr[i] = offset == -1;
+            if (offset != -1){{
+                value_data[offset] = value_ptr[i];
+            }}
+        }}
+        """)
+        return code 
 
 class HashTable(pccm.Class, pccm.pybind.PybindClassMixin):
     """a simple hashtable for both cpu and cuda.
@@ -451,6 +478,76 @@ class HashTable(pccm.Class, pccm.pybind.PybindClassMixin):
             """)
         return code 
 
+    @pccm.pybind.mark 
+    @_member_func
+    def insert_exist_keys(self):
+        """insert v of given k if k exists. won't insert any new key.
+        """
+        code = pccm.FunctionCode()
+        if not CUMM_CPU_ONLY_BUILD:
+            code.add_dependency(TensorViewHashKernel, HashTableKernel)
+        code.arg("keys", "tv::Tensor")
+        code.arg("values", "tv::Tensor")
+        code.arg("is_empty", "tv::Tensor")
+
+        code.arg("stream", "std::uintptr_t")
+
+        code.raw(f"""
+        auto N = keys.dim(0);
+        TV_ASSERT_RT_ERR(keys.itemsize() == key_itemsize_, "keys itemsize not equal to", key_itemsize_);
+        TV_ASSERT_RT_ERR(values.itemsize() == value_itemsize_, "values itemsize not equal to", value_itemsize_);
+        TV_ASSERT_RT_ERR(N == values.dim(0) && is_empty.dim(0) == N, "number of key and value must same");
+        auto is_empty_ptr = is_empty.data_ptr<uint8_t>();
+        """)
+        with code.if_("is_cpu"):
+            map_name = "cpu_map"
+            # here it's safe to use omp in query.
+            for k_type, v_type in self.cpu_map_storage_select("key_itemsize_", "value_itemsize_", map_name, code):
+                code.raw(f"""
+                auto k_ptr = reinterpret_cast<{k_type}*>(keys.raw_data());
+                auto v_ptr = reinterpret_cast<{v_type}*>(values.raw_data());
+                tv::kernel_1d_cpu(keys.device(), N, [&](size_t begin, size_t end, size_t step){{
+                    bool emp;
+                    for (size_t i = begin; i < end; i += step){{
+                        auto iter = {map_name}.find(k_ptr[i]);
+                        emp = iter == {map_name}.end();
+                        if (!emp){{
+                            iter.value() = v_ptr[i];
+                        }}
+                        is_empty_ptr[i] = uint8_t(emp);
+                    }}
+                }});
+                """)
+        if not CUMM_CPU_ONLY_BUILD:
+            with code.else_():
+                code.raw(f"""
+                auto custream = reinterpret_cast<cudaStream_t>(stream);
+                """)
+                for k_items in _dispatch_ints(code, [4, 8], "keys_data.itemsize()"):
+                    code.raw(f"""
+                    using K = tv::hash::itemsize_to_unsigned_t<{k_items}>;
+                    constexpr K kEmptyKey = std::numeric_limits<K>::max();
+                    K* key_data_ptr = reinterpret_cast<K*>(keys_data.raw_data());
+                    const K* key_ptr = reinterpret_cast<const K*>(keys.raw_data());
+
+                    """)
+                    for v_items in _dispatch_ints(code, [4, 8], "values_data.itemsize()"):
+                        code.raw(f"""
+                        using V = tv::hash::itemsize_to_unsigned_t<{v_items}>;
+                        V* value_data_ptr = reinterpret_cast<V*>(values_data.raw_data());
+                        const V* value_ptr = reinterpret_cast<const V*>(values.raw_data());
+                        using table_t =
+                            tv::hash::LinearHashTableSplit<K, V, tv::hash::Murmur3Hash<K>,
+                                                        kEmptyKey, false>;
+                        table_t table(key_data_ptr, value_data_ptr, keys_data.dim(0));
+                        tv::cuda::Launch launcher(N, custream);
+                        launcher(insert_exist_keys_kernel<table_t>, table, key_ptr, value_ptr, is_empty_ptr, size_t(N));
+                        """)
+        else:
+            code.raw(f"""
+            TV_THROW_RT_ERR("spconv not compiled with cuda, don't support cuda");
+            """)
+        return code 
 
     def cpu_map_storage_select(self, k_itemsize: str, v_itemsize: str, res_var: str, code: pccm.FunctionCode):
         different_kvs = [(4, 4), (4, 8), (8, 4), (8, 8)]

spconv/pytorch/conv.py

@@ -29,7 +29,7 @@ from spconv.debug_utils import spconv_save_debug_data
 from spconv.pytorch import functional as Fsp
 from spconv.pytorch import ops
 from spconv.cppconstants import CPU_ONLY_BUILD
-from spconv.pytorch.core import IndiceData, SparseConvTensor, ImplicitGemmIndiceData
+from spconv.pytorch.core import IndiceData, SparseConvTensor, ImplicitGemmIndiceData, expand_nd
 from spconv.pytorch.modules import SparseModule
 from spconv.constants import FILTER_HWIO
 from spconv.utils import nullcontext
@@ -95,32 +95,22 @@ class SparseConvolution(SparseModule):
                  name=None):
         super(SparseConvolution, self).__init__(name=name)
         assert groups == 1, "don't support groups for now"
-        if not isinstance(kernel_size, (list, tuple)):
-            kernel_size = [kernel_size] * ndim
-        if not isinstance(stride, (list, tuple)):
-            stride = [stride] * ndim
-        if not isinstance(padding, (list, tuple)):
-            padding = [padding] * ndim
-        if not isinstance(dilation, (list, tuple)):
-            dilation = [dilation] * ndim
-        if not isinstance(output_padding, (list, tuple)):
-            output_padding = [output_padding] * ndim
         self.ndim = ndim
         self.in_channels = in_channels
         self.out_channels = out_channels
-        self.kernel_size = kernel_size
+        self.kernel_size = expand_nd(ndim, kernel_size)
         kv = int(np.prod(kernel_size))
         kv_stride = int(np.prod(stride))
         self.conv1x1 = kv == 1
         # TODO we should deprecate support for ksize == 1 but stride != 1.
         if not subm:
             self.conv1x1 &= kv_stride == 1
-        self.stride = stride
-        self.padding = padding
-        self.dilation = dilation
+        self.stride = expand_nd(ndim, stride)
+        self.padding = expand_nd(ndim, padding)
+        self.dilation = expand_nd(ndim, dilation)
         self.transposed = transposed
         self.inverse = inverse
-        self.output_padding = output_padding
+        self.output_padding = expand_nd(ndim, output_padding)
         self.groups = groups
         self.subm = subm
         self.indice_key = indice_key
@@ -142,15 +132,15 @@ class SparseConvolution(SparseModule):
             if FILTER_HWIO:
                 # RSCK
                 self.weight = Parameter(
-                    torch.Tensor(*kernel_size, in_channels, out_channels))
+                    torch.Tensor(*self.kernel_size, in_channels, out_channels))
             else:
                 # RSKC
                 self.weight = Parameter(
-                    torch.Tensor(*kernel_size, out_channels, in_channels))
+                    torch.Tensor(*self.kernel_size, out_channels, in_channels))
         else:
             # KRSC
             self.weight = Parameter(
-                torch.Tensor(out_channels, *kernel_size, in_channels))
+                torch.Tensor(out_channels, *self.kernel_size, in_channels))
 
         if bias:
             self.bias = Parameter(torch.Tensor(out_channels))
@@ -321,7 +311,11 @@ class SparseConvolution(SparseModule):
                                                  spatial_shape,
                                                  out_spatial_shape,
                                                  is_subm=self.subm,
-                                                 algo=algo)
+                                                 algo=algo,
+                                                 ksize=self.kernel_size,
+                                                 stride=self.stride,
+                                                 padding=self.padding,
+                                                 dilation=self.dilation)
                         if self.indice_key is not None:
                             msg = f"your indice key {self.indice_key} already exists in this sparse tensor."
                             assert self.indice_key not in indice_dict, msg
@@ -364,10 +358,7 @@ class SparseConvolution(SparseModule):
                     mask_argsort_bwd_splits = datas.mask_argsort_fwd_splits
                     masks = datas.masks
                     out_spatial_shape = datas.spatial_shape
-                    assert datas.pair_fwd.shape[0] == np.prod(
-                        self.kernel_size
-                    ), "inverse conv must have same kernel size as its couple conv"
-
+                    assert datas.ksize == self.kernel_size, "inverse conv must have same kernel size as its couple conv"
                 else:
                     if self.indice_key is not None and datas is not None:
                         outids = datas.out_indices
@@ -378,6 +369,19 @@ class SparseConvolution(SparseModule):
                         mask_argsort_fwd_splits = datas.mask_argsort_fwd_splits
                         mask_argsort_bwd_splits = datas.mask_argsort_bwd_splits
                         masks = datas.masks
+                        assert datas.is_subm, "only support reuse subm indices"
+                        if self.kernel_size != datas.ksize:
+                            raise ValueError(f"subm with same indice_key must have same kernel"
+                                f" size, expect {datas.ksize}, this layer {self.kernel_size}")
+                        if self.dilation != datas.dilation:
+                            raise ValueError(f"subm with same indice_key must have same dilation"
+                                f", expect {datas.dilation}, this layer {self.dilation}")
+                        if input.spatial_shape != datas.spatial_shape:
+                            raise ValueError(f"subm with same indice_key must have same spatial structure"
+                                f", expect {datas.spatial_shape}, input {spatial_shape}")
+                        if input.indices.shape[0] != datas.indices.shape[0]:
+                            raise ValueError(f"subm with same indice_key must have same num of indices"
+                                f", expect {datas.indices.shape[0]}, input {input.indices.shape[0]}")
                     else:
 
                         with input._timer.namespace("gen_pairs"):
@@ -432,7 +436,11 @@ class SparseConvolution(SparseModule):
                                 is_subm=self.subm,
                                 spatial_shape=spatial_shape,
                                 out_spatial_shape=out_spatial_shape,
-                                algo=algo)
+                                algo=algo,
+                                ksize=self.kernel_size,
+                                stride=self.stride,
+                                padding=self.padding,
+                                dilation=self.dilation)
                             msg = f"your indice key {self.indice_key} already exists in this sparse tensor."
                             assert self.indice_key not in indice_dict, msg
                             indice_dict[self.indice_key] = indice_data

spconv/pytorch/core.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional, Union
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -58,7 +58,8 @@ class ThrustSortAllocator:
 
 class IndiceData(object):
     def __init__(self, out_indices, indices, indice_pairs, indice_pair_num,
-                 spatial_shape, out_spatial_shape, is_subm: bool, algo: ConvAlgo):
+                 spatial_shape, out_spatial_shape, is_subm: bool, algo: ConvAlgo,
+                 ksize: List[int], stride: List[int], dilation: List[int], padding: List[int]):
         self.out_indices = out_indices
         self.indices = indices
         self.indice_pairs = indice_pairs
@@ -67,6 +68,10 @@ class IndiceData(object):
         self.out_spatial_shape = out_spatial_shape
         self.is_subm = is_subm
         self.algo = algo
+        self.ksize = ksize
+        self.stride = stride
+        self.dilation = dilation
+        self.padding = padding
 
 
 class ImplicitGemmIndiceData(object):
@@ -77,7 +82,8 @@ class ImplicitGemmIndiceData(object):
                  mask_argsort_fwd_splits: List[torch.Tensor],
                  mask_argsort_bwd_splits: List[torch.Tensor],
                  masks: List[np.ndarray], spatial_shape, 
-                 out_spatial_shape, is_subm: bool, algo: ConvAlgo):
+                 out_spatial_shape, is_subm: bool, algo: ConvAlgo,
+                 ksize: List[int], stride: List[int], dilation: List[int], padding: List[int]):
         self.out_indices = out_indices
         self.indices = indices
         self.pair_fwd = pair_fwd
@@ -91,6 +97,10 @@ class ImplicitGemmIndiceData(object):
         self.out_spatial_shape = out_spatial_shape
         self.is_subm = is_subm
         self.algo = algo
+        self.ksize = ksize
+        self.stride = stride
+        self.dilation = dilation
+        self.padding = padding
 
 
 def scatter_nd(indices, updates, shape):
@@ -235,3 +245,13 @@ class SparseConvTensor(metaclass=SpConvTensorMeta):
         tensor.thrust_allocator = self.thrust_allocator
         tensor._timer = self._timer
         return tensor
+
+def expand_nd(ndim: int, val: Union[int, List[int], Tuple[int, ...]]) -> List[int]:
+    if isinstance(val, int):
+        res = [val] * ndim 
+    elif isinstance(val, tuple):
+        res = list(val)
+    else:
+        res = val
+    assert len(res) == ndim
+    return res 

spconv/pytorch/functional.py

@@ -29,6 +29,7 @@ from pathlib import Path
 from spconv.pytorch.hash import HashTable
 from cumm.gemm.layout import to_stride
 from typing import List
+from functools import reduce 
 
 _MAX_INT32 = 2147483647
 
@@ -365,7 +366,7 @@ indice_maxpool_implicit_gemm = SparseMaxPoolImplicitGemmFunction.apply
 def _indice_to_scalar(indices: torch.Tensor, shape: List[int]):
     assert indices.shape[1] == len(shape)
     stride = to_stride(np.array(shape, dtype=np.int64))
-    scalar_inds = indices[:, -1]
+    scalar_inds = indices[:, -1].clone()
     for i in range(len(shape) - 1):
         scalar_inds += stride[i] * indices[:, i]
     return scalar_inds.contiguous()
@@ -426,17 +427,34 @@ def sparse_add_hash_based(*tens: SparseConvTensor):
     res.thrust_allocator = first.thrust_allocator
     return res 
 
-def sparse_add(a: SparseConvTensor, b: SparseConvTensor):
-    assert a.spatial_shape == b.spatial_shape
-    assert a.batch_size == b.batch_size
-    assert a.features.shape[1] == a.features.shape[1]
-    res_shape = [a.batch_size, *a.spatial_shape, a.features.shape[1]]
-
-    a_th = torch.sparse_coo_tensor(a.indices.T, a.features, res_shape, requires_grad=True)
-    b_th = torch.sparse_coo_tensor(b.indices.T, b.features, res_shape, requires_grad=True)
+def sparse_add(*tens: SparseConvTensor):
+    """reuse torch.sparse. the internal is sort + unique 
+    """
+    max_num_indices = 0
+    max_num_indices_idx = 0
+    ten_ths: List[torch.Tensor] = []
+    first = tens[0]
+    res_shape = [first.batch_size, *first.spatial_shape, first.features.shape[1]]
 
-    c_th = (a_th + b_th).coalesce()
+    for i, ten in enumerate(tens):
+        assert ten.spatial_shape == tens[0].spatial_shape
+        assert ten.batch_size == tens[0].batch_size
+        assert ten.features.shape[1] == tens[0].features.shape[1]
+        if max_num_indices < ten.features.shape[0]:
+            max_num_indices_idx = i
+            max_num_indices = ten.features.shape[0]
+        ten_ths.append(torch.sparse_coo_tensor(ten.indices.T, ten.features, res_shape, requires_grad=True))
+    
+    c_th = reduce(lambda x, y: x + y, ten_ths).coalesce()
     c_th_inds = c_th.indices().T.contiguous().int()
     c_th_values = c_th.values()
     assert c_th_values.is_contiguous()
-    return SparseConvTensor(c_th_values, c_th_inds, a.spatial_shape, a.batch_size)
+
+    res = SparseConvTensor(c_th_values, c_th_inds, first.spatial_shape, first.batch_size, 
+        benchmark=first.benchmark)
+    if c_th_values.shape[0] == max_num_indices:
+        res.indice_dict = tens[max_num_indices_idx].indice_dict
+    res.benchmark_record = first.benchmark_record
+    res._timer = first._timer 
+    res.thrust_allocator = first.thrust_allocator
+    return res 

spconv/pytorch/hash.py

@@ -78,6 +78,9 @@ class HashTable:
         return self._table.insert(keys_tv, values_tv, stream)
 
     def query(self, keys: torch.Tensor, values: Optional[torch.Tensor] = None):
+        """query value by keys, if values is not None, create a new one.
+        return values and a uint8 tensor that whether query success.
+        """
         keys_tv = torch_tensor_to_tv(keys)
         if values is None:
             values = torch.empty([keys.shape[0]], dtype=self.value_dtype, device=keys.device)
@@ -90,7 +93,24 @@ class HashTable:
         self._table.query(keys_tv, values_tv, is_empty_tv, stream)
         return values, is_empty
 
+    def insert_exist_keys(self, keys: torch.Tensor, values: torch.Tensor):
+        """insert kv that k exists in table. return a uint8 tensor that
+        whether insert success.
+        """
+        keys_tv = torch_tensor_to_tv(keys)
+        values_tv = torch_tensor_to_tv(values)
+        stream = 0
+        if not self.is_cpu:
+            stream = get_current_stream()
+        is_success = torch.empty([keys.shape[0]], dtype=torch.uint8, device=keys.device)
+        is_success_tv = torch_tensor_to_tv(is_success)
+        self._table.insert_exist_keys(keys_tv, values_tv, is_success_tv, stream)
+        return is_success
+
     def assign_arange_(self):
+        """iterate table, assign values with "arange" value.
+        equivalent to 1. get key by items(), 2. use key and arange(key.shape[0]) to insert
+        """
         count_tv = tv.Tensor()
         count = torch.Tensor()
         stream = 0

spconv/pytorch/pool.py

@@ -26,7 +26,7 @@ from spconv import pytorch as spconv
 from spconv.core import ConvAlgo
 from spconv.pytorch import functional as Fsp
 from spconv.pytorch import ops
-from spconv.pytorch.core import IndiceData, ImplicitGemmIndiceData
+from spconv.pytorch.core import IndiceData, ImplicitGemmIndiceData, expand_nd
 from spconv.pytorch.modules import SparseModule
 from spconv.cppconstants import CPU_ONLY_BUILD
 from spconv.utils import nullcontext
@@ -36,7 +36,7 @@ class SparseMaxPool(SparseModule):
     def __init__(self,
                  ndim,
                  kernel_size: Union[int, List[int], Tuple[int, ...]] = 3,
-                 stride: Union[int, List[int], Tuple[int, ...]] = 1,
+                 stride: Optional[Union[int, List[int], Tuple[int, ...]]] = 1,
                  padding: Union[int, List[int], Tuple[int, ...]] = 0,
                  dilation: Union[int, List[int], Tuple[int, ...]] = 1,
                  indice_key: Optional[str] = None,
@@ -44,22 +44,15 @@ class SparseMaxPool(SparseModule):
                  algo: Optional[ConvAlgo] = None,
                  name=None):
         super(SparseMaxPool, self).__init__(name=name)
-        if not isinstance(kernel_size, (list, tuple)):
-            kernel_size = [kernel_size] * ndim
-        if stride is None:
-            stride = kernel_size.copy()
-        if not isinstance(stride, (list, tuple)):
-            stride = [stride] * ndim
-        if not isinstance(padding, (list, tuple)):
-            padding = [padding] * ndim
-        if not isinstance(dilation, (list, tuple)):
-            dilation = [dilation] * ndim
         self.ndim = ndim
-        self.kernel_size = kernel_size
-        self.stride = stride
-        self.padding = padding
+        self.kernel_size = expand_nd(ndim, kernel_size)
+        if stride is None:
+            self.stride = self.kernel_size.copy()
+        else:
+            self.stride = expand_nd(ndim, stride)
+        self.padding = expand_nd(ndim, padding)
         self.subm = subm
-        self.dilation = dilation
+        self.dilation = expand_nd(ndim, dilation)
         self.indice_key = indice_key
         kv = int(np.prod(kernel_size))
         if algo is None:
@@ -155,7 +148,11 @@ class SparseMaxPool(SparseModule):
                                                  spatial_shape,
                                                  out_spatial_shape,
                                                  is_subm=False,
-                                                 algo=self.algo)
+                                                 algo=self.algo,
+                                                 ksize=self.kernel_size,
+                                                 stride=self.stride,
+                                                 padding=self.padding,
+                                                 dilation=self.dilation)
                         indice_dict[self.indice_key] = indice_data
                     else:
                         raise ValueError(
@@ -204,7 +201,11 @@ class SparseMaxPool(SparseModule):
                         is_subm=self.subm,
                         spatial_shape=spatial_shape,
                         out_spatial_shape=out_spatial_shape,
-                        algo=self.algo)
+                        algo=self.algo,
+                        ksize=self.kernel_size,
+                        stride=self.stride,
+                        padding=self.padding,
+                        dilation=self.dilation)
                     msg = f"your indice key {self.indice_key} already exists in this sparse tensor."
                     assert self.indice_key not in indice_dict, msg
                     indice_dict[self.indice_key] = indice_data

version.txt

-2.1.16
\ No newline at end of file
+2.1.17
\ No newline at end of file