tensor input support for sparse_quantize (fix NVIDIA#389)

MinkowskiEngine/utils/quantization.py

@@ -122,6 +122,17 @@ def quantize_label(coords, labels, ignore_label):
         return MEB.quantize_label_th(coords, labels.int(), ignore_label)
 
 
+def _auto_floor(array):
+    assert isinstance(
+        array, (np.ndarray, torch.Tensor)
+    ), "array must be either np.array or torch.Tensor."
+
+    if isinstance(array, np.ndarray):
+        return np.floor(array)
+    else:
+        return torch.floor(array)
+
+
 def sparse_quantize(
     coordinates,
     features=None,
@@ -235,23 +246,21 @@ def sparse_quantize(
             ), "Quantization size and coordinates size mismatch."
             if isinstance(coordinates, np.ndarray):
                 quantization_size = np.array([i for i in quantization_size])
-                discrete_coordinates = np.floor(coordinates / quantization_size)
             else:
                 quantization_size = torch.Tensor([i for i in quantization_size])
-                discrete_coordinates = (coordinates / quantization_size).floor()
+            discrete_coordinates = _auto_floor(coordinates / quantization_size)
 
         elif np.isscalar(quantization_size):  # Assume that it is a scalar
 
             if quantization_size == 1:
-                discrete_coordinates = coordinates
+                discrete_coordinates = _auto_floor(coordinates)
             else:
-                discrete_coordinates = np.floor(coordinates / quantization_size)
+                discrete_coordinates = _auto_floor(coordinates / quantization_size)
         else:
             raise ValueError("Not supported type for quantization_size.")
     else:
-        discrete_coordinates = coordinates
+        discrete_coordinates = _auto_floor(coordinates)
 
-    discrete_coordinates = np.floor(discrete_coordinates)
     if isinstance(coordinates, np.ndarray):
         discrete_coordinates = discrete_coordinates.astype(np.int32)
     else: