Skip to content

GitLab

Commit cf15a2dd authored by rusty1s's avatar rusty1s
Browse files

c99 c

parent af085f49
Show whitespace changes
Inline Side-by-side
Showing with 12 additions and 6 deletions
torch_scatter/src/generic/cpu.c
View file @ cf15a2dd
...@@ -3,24 +3,27 @@ ...@@ -3,24 +3,27 @@
#else #else
void scatter_(mul)(int dim, THTensor *output, THLongTensor *index, THTensor *input) { void scatter_(mul)(int dim, THTensor *output, THLongTensor *index, THTensor *input) {
int64_t i;
TH_TENSOR_DIM_APPLY3(real, output, int64_t, index, real, input, dim, TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM, TH_TENSOR_DIM_APPLY3(real, output, int64_t, index, real, input, dim, TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM,
for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) { for (i = 0; i < THLongTensor_size(index, dim); i++) {
assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter); assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
output_data[index_data[i]] *= input_data[i]; output_data[index_data[i]] *= input_data[i];
}) })
} }
void scatter_(div)(int dim, THTensor *output, THLongTensor *index, THTensor *input) { void scatter_(div)(int dim, THTensor *output, THLongTensor *index, THTensor *input) {
int64_t i;
TH_TENSOR_DIM_APPLY3(real, output, int64_t, index, real, input, dim, TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM, TH_TENSOR_DIM_APPLY3(real, output, int64_t, index, real, input, dim, TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM,
for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) { for (i = 0; i < THLongTensor_size(index, dim); i++) {
assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter); assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
output_data[index_data[i]] /= input_data[i]; output_data[index_data[i]] /= input_data[i];
}) })
} }
void scatter_(mean)(int dim, THTensor *output, THLongTensor *index, THTensor *input, THTensor *count) { void scatter_(mean)(int dim, THTensor *output, THLongTensor *index, THTensor *input, THTensor *count) {
int64_t i;
TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, input, real, count, dim, TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, input, real, count, dim,
for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) { for (i = 0; i < THLongTensor_size(index, dim); i++) {
assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter); assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
output_data[index_data[i]] += input_data[i]; output_data[index_data[i]] += input_data[i];
count_data[index_data[i]]++; count_data[index_data[i]]++;
...@@ -28,8 +31,9 @@ void scatter_(mean)(int dim, THTensor *output, THLongTensor *index, THTensor *in ...@@ -28,8 +31,9 @@ void scatter_(mean)(int dim, THTensor *output, THLongTensor *index, THTensor *in
} }
void scatter_(max)(int dim, THTensor *output, THLongTensor *index, THTensor *input, THLongTensor *arg) { void scatter_(max)(int dim, THTensor *output, THLongTensor *index, THTensor *input, THLongTensor *arg) {
int64_t i;
TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, input, int64_t, arg, dim, TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, input, int64_t, arg, dim,
for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) { for (i = 0; i < THLongTensor_size(index, dim); i++) {
assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter); assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
if (input_data[i] >= output_data[index_data[i]]) { if (input_data[i] >= output_data[index_data[i]]) {
output_data[index_data[i]] = input_data[i]; output_data[index_data[i]] = input_data[i];
...@@ -39,8 +43,9 @@ void scatter_(max)(int dim, THTensor *output, THLongTensor *index, THTensor *inp ...@@ -39,8 +43,9 @@ void scatter_(max)(int dim, THTensor *output, THLongTensor *index, THTensor *inp
} }
void scatter_(min)(int dim, THTensor *output, THLongTensor *index, THTensor *input, THLongTensor *arg) { void scatter_(min)(int dim, THTensor *output, THLongTensor *index, THTensor *input, THLongTensor *arg) {
int64_t i;
TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, input, int64_t, arg, dim, TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, input, int64_t, arg, dim,
for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) { for (i = 0; i < THLongTensor_size(index, dim); i++) {
assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter); assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
if (input_data[i] <= output_data[index_data[i]]) { if (input_data[i] <= output_data[index_data[i]]) {
output_data[index_data[i]] = input_data[i]; output_data[index_data[i]] = input_data[i];
...@@ -50,8 +55,9 @@ void scatter_(min)(int dim, THTensor *output, THLongTensor *index, THTensor *inp ...@@ -50,8 +55,9 @@ void scatter_(min)(int dim, THTensor *output, THLongTensor *index, THTensor *inp
} }
void index_backward(int dim, THTensor *output, THLongTensor *index, THTensor *grad, THLongTensor *arg) { void index_backward(int dim, THTensor *output, THLongTensor *index, THTensor *grad, THLongTensor *arg) {
int64_t i;
TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, grad, int64_t, arg, dim, TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, grad, int64_t, arg, dim,
for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) { for (i = 0; i < THLongTensor_size(index, dim); i++) {
if (arg_data[index_data[i]] == i) output_data[i] = grad_data[index_data[i]]; if (arg_data[index_data[i]] == i) output_data[i] = grad_data[index_data[i]];
}) })
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment