none arguments

Author: rusty1s

aten/cpu/cluster.cpp

@@ -7,9 +7,6 @@ at::Tensor graclus(at::Tensor row, at::Tensor col, at::Tensor weight) {
 
 
 at::Tensor grid(at::Tensor pos, at::Tensor size, at::Tensor start, at::Tensor end) {
-  if (!start.defined()) start = std::get<0>(pos.min(1));
-  if (!end.defined()) end = std::get<0>(pos.max(1));
-
   size = size.toType(pos.type());
   start = start.toType(pos.type());
   end = end.toType(pos.type());
@@ -30,7 +27,6 @@ at::Tensor grid(at::Tensor pos, at::Tensor size, at::Tensor start, at::Tensor en
 
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("graclus", &graclus, "Graclus (CPU)", py::arg("row"), py::arg("col"), py::arg("weight"));
-  m.def("grid", &grid, "Grid (CPU)", py::arg("pos"), py::arg("size"), py::arg("start"),
-        py::arg("end"));
+  m.def("graclus", &graclus, "Graclus (CPU)");
+  m.def("grid", &grid, "Grid (CPU)");
 }

aten/cpu/cluster.py

@@ -3,15 +3,17 @@
 import cluster_cpu
 
 
-def grid_cluster(pos, size, start, end):
+def grid_cluster(pos, size, start=None, end=None):
+    start = pos.t().min(dim=1)[0] if start is None else start
+    end = pos.t().max(dim=1)[0] if end is None else end
     return cluster_cpu.grid(pos, size, start, end)
 
 
-pos = torch.tensor([[1, 1], [3, 3], [5, 5], [7, 7]], dtype=torch.uint8)
+pos = torch.tensor([[1, 1], [3, 3], [5, 5], [7, 7]])
 size = torch.tensor([2, 2])
 start = torch.tensor([0, 0])
 end = torch.tensor([7, 7])
 print('pos', pos.tolist())
 print('size', size.tolist())
-cluster = grid_cluster(pos, size, start, end)
+cluster = grid_cluster(pos, size)
 print('result', cluster.tolist(), cluster.dtype)