clean up

aten/TH/THGreedy.c

 #include <TH/TH.h>
 
 #define THGreedy_(NAME) TH_CONCAT_4(TH,Real,Greedy_,NAME)
+#define DATA(TENSOR) TENSOR->storage->data + TENSOR->storageOffset
 
 #define TH_GREEDY_CLUSTER(cluster, row, col, deg, SELECT) { \
-  int64_t *clusterData = cluster->storage->data + cluster->storageOffset; \
-  int64_t *rowData = row->storage->data + row->storageOffset; \
-  int64_t *colData = col->storage->data + col->storageOffset; \
-  int64_t *degData = deg->storage->data + deg->storageOffset; \
+  THLongTensor_fill(cluster, -1); \
+  int64_t *clusterData = DATA(cluster); \
+  int64_t *rowData = DATA(row); \
+  int64_t *colData = DATA(col); \
+  int64_t *degData = DATA(deg); \
   ptrdiff_t rowIdx = 0, neighborIdx; \
   int64_t rowValue, colValue, clusterValue, tmp; \
   while(rowIdx < THLongTensor_nElement(row)) { \
     rowValue = rowData[rowIdx]; \
+    printf("rowValue = %lli, ", rowValue); \
     if (clusterData[rowValue] < 0) { \
       colValue = rowValue; \
       SELECT \
       clusterValue = rowValue < colValue ? rowValue : colValue; \
+      printf("%lli", clusterValue); \
       clusterData[rowValue] = clusterValue; \
       clusterData[colValue] = clusterValue; \
     } \

aten/TH/THGrid.c

 #include <TH/TH.h>
 
 #define THGrid_(NAME) TH_CONCAT_4(TH,Real,Grid_,NAME)
+#define DATA(TENSOR) TENSOR->storage->data + TENSOR->storageOffset
 
 #include "generic/THGrid.c"
 #include "THGenerateAllTypes.h"

aten/TH/THGrid.h

-void THByteGrid(THLongTensor *cluster,
+void THByteGrid_cluster(THLongTensor *cluster,
                         THByteTensor *pos,
                         THByteTensor *size,
                         THLongTensor *count);
 
-void THCharGrid(THLongTensor *cluster,
+void THCharGrid_cluster(THLongTensor *cluster,
                         THCharTensor *pos,
                         THCharTensor *size,
                         THLongTensor *count);
 
-void THShortGrid(THLongTensor *cluster,
+void THShortGrid_cluster(THLongTensor *cluster,
                          THShortTensor *pos,
                          THShortTensor *size,
                          THLongTensor *count);
 
-void THIntGrid(THLongTensor *cluster,
+void THIntGrid_cluster(THLongTensor *cluster,
                        THIntTensor *pos,
                        THIntTensor *size,
                        THLongTensor *count);
 
-void THLongGrid(THLongTensor *cluster,
+void THLongGrid_cluster(THLongTensor *cluster,
                         THLongTensor *pos,
                         THLongTensor *size,
                         THLongTensor *count);
 
-void THFloatGrid(THLongTensor *cluster,
+void THFloatGrid_cluster(THLongTensor *cluster,
                          THFloatTensor *pos,
                          THFloatTensor *size,
                          THLongTensor *count);
 
-void THDoubleGrid(THLongTensor *cluster,
+void THDoubleGrid_cluster(THLongTensor *cluster,
                           THDoubleTensor *pos,
                           THDoubleTensor *size,
                           THLongTensor *count);

aten/TH/generic/THGreedy.c

@@ -4,7 +4,7 @@
 
 void THGreedy_(cluster)(THLongTensor *cluster, THLongTensor *row, THLongTensor *col,
                         THLongTensor *deg, THTensor *weight) {
-  real *weightData = weight->storage->data + weight->storageOffset;
+  real *weightData = DATA(weight);
   real maxWeight = 0, tmpWeight;
   TH_GREEDY_CLUSTER(cluster, row, col, deg,
     for (neighborIdx = rowIdx; neighborIdx < rowIdx + degData[rowValue]; neighborIdx++) {

aten/TH/generic/THGrid.c

@@ -3,20 +3,21 @@
 #else
 
 void THGrid_(cluster)(THLongTensor *cluster, THTensor *pos, THTensor *size, THLongTensor *count) {
-  real *sizeData = size->storage->data + size->storageOffset;
-  int64_t *countData = count->storage->data + count->storageOffset;
-  int64_t dims = THLongTensor_nElement(count);
-  THLongTensor_unsqueeze1d(cluster, NULL, 1);
-  ptrdiff_t d; int64_t coef, value;
-  TH_TENSOR_DIM_APPLY2(int64_t, cluster, real, pos, 1,
+  int64_t *clusterData = DATA(cluster);
+  real *posData = DATA(pos);
+  real *sizeData = DATA(size);
+  int64_t *countData = DATA(count);
+
+  ptrdiff_t n, d; int64_t coef, value;
+  for (n = 0; n < THTensor_(size)(pos, 0); n++) {
     coef = 1; value = 0;
-    for (d = 0; d < dims; d++) {
-      value += coef * (int64_t) (*(pos_data + d * pos_stride) / sizeData[d]);
+    for (d = 0; d < THTensor_(size)(pos, 1); d++) {
+      value += coef * (int64_t) (*(posData + d * pos->stride[1]) / sizeData[d]);
       coef *= countData[d];
     }
-    cluster_data[0] = value;
-  )
-  THLongTensor_squeeze1d(cluster, NULL, 1);
+    posData += pos->stride[0];
+    clusterData[n] = value;
+  }
 }
 
 #endif  // TH_GENERIC_FILE