arg check

torch_scatter/kernel/generic/kernel.cu

@@ -2,31 +2,39 @@
 #define THC_GENERIC_FILE "generic/kernel.cu"
 #else
 
-void scatter_(mul)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input) {
+void check(THCState *state, THCTensor *output, THCudaLongTensor *index, THCTensor *input) {
   THCAssertSameGPU(THCTensor_(checkGPU)(state, 1, output, input));
   THCAssertSameGPU(THCudaLongTensor_checkGPU(state, 2, index));
   THArgCheck(THCTensor_(nDimension)(state, output) <= MAX_DIMS, 1, "Tensor too large or too many dimensions");
+}
 
+void scatter_(mul)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input) {
+  check(state, output, index, input);
   printf("mul");
 }
 
 void scatter_(div)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input) {
+  check(state, output, index, input);
   printf("div");
 }
 
 void scatter_(mean)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input, THCTensor *num_output) {
+  check(state, output, index, input);
   printf("mean");
 }
 
 void scatter_(max)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input, THCudaLongTensor *arg_output) {
+  check(state, output, index, input);
   printf("max");
 }
 
 void scatter_(min)(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *input, THCudaLongTensor *arg_output) {
+  check(state, output, index, input);
   printf("min");
 }
 
 void index_backward(THCState *state, int dim, THCTensor *output, THCudaLongTensor *index, THCTensor *grad, THCudaLongTensor *arg_grad) {
+  check(state, output, index, grad);
   printf("index_backward");
 }
 

torch_scatter/kernel/kernel.cu

@@ -4,6 +4,7 @@
 
 #define scatter_(NAME) TH_CONCAT_4(scatter_, NAME, _kernel_, Real)
 #define index_backward TH_CONCAT_2(index_backward_kernel_, Real)
+#define check TH_CONCAT_2(check_kernel_, Real)
 
 #define MAX_DIMS 25