cpu.c 2.72 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
#ifndef TH_GENERIC_FILE
rusty1s's avatar
rusty1s committed
2
#define TH_GENERIC_FILE "generic/cpu.c"
rusty1s's avatar
rusty1s committed
3
4
#else

rusty1s's avatar
rusty1s committed
5
6
void scatter_(mul)(int dim, THTensor *output, THLongTensor *index, THTensor *input) {
  TH_TENSOR_DIM_APPLY3(real, output, int64_t, index, real, input, dim, TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM,
rusty1s's avatar
rusty1s committed
7
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
rusty1s's avatar
rusty1s committed
8
9
      assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
      output_data[index_data[i]] *= input_data[i];
rusty1s's avatar
rusty1s committed
10
11
12
    })
}

rusty1s's avatar
rusty1s committed
13
14
void scatter_(div)(int dim, THTensor *output, THLongTensor *index, THTensor *input) {
  TH_TENSOR_DIM_APPLY3(real, output, int64_t, index, real, input, dim, TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM,
rusty1s's avatar
rusty1s committed
15
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
rusty1s's avatar
rusty1s committed
16
17
      assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
      output_data[index_data[i]] /= input_data[i];
rusty1s's avatar
rusty1s committed
18
19
20
    })
}

rusty1s's avatar
rusty1s committed
21
void scatter_(mean)(int dim, THTensor *output, THLongTensor *index, THTensor *input, THTensor *output_count) {
rusty1s's avatar
rusty1s committed
22
  TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, input, real, output_count, dim,
rusty1s's avatar
rusty1s committed
23
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
rusty1s's avatar
rusty1s committed
24
25
26
      assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
      output_data[index_data[i]] += input_data[i];
      output_count_data[index_data[i]]++;
rusty1s's avatar
rusty1s committed
27
28
29
    })
}

rusty1s's avatar
renames  
rusty1s committed
30
31
void scatter_(max)(int dim, THTensor *output, THLongTensor *index, THTensor *input, THLongTensor *output_arg) {
  TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, input, int64_t, output_arg, dim,
rusty1s's avatar
min max  
rusty1s committed
32
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
rusty1s's avatar
rusty1s committed
33
34
35
      assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
      if (input_data[i] >= output_data[index_data[i]]) {
        output_data[index_data[i]] = input_data[i];
rusty1s's avatar
renames  
rusty1s committed
36
        output_arg_data[index_data[i]] = i;
rusty1s's avatar
rusty1s committed
37
      }
rusty1s's avatar
min max  
rusty1s committed
38
39
40
    })
}

rusty1s's avatar
renames  
rusty1s committed
41
42
void scatter_(min)(int dim, THTensor *output, THLongTensor *index, THTensor *input, THLongTensor *output_arg) {
  TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, input, int64_t, output_arg, dim,
rusty1s's avatar
min max  
rusty1s committed
43
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
rusty1s's avatar
rusty1s committed
44
45
46
      assertIndexInBoundaries(index_data[i], output_size, TH_TENSOR_DIM_APPLY_counter);
      if (input_data[i] <= output_data[index_data[i]]) {
        output_data[index_data[i]] = input_data[i];
rusty1s's avatar
renames  
rusty1s committed
47
        output_arg_data[index_data[i]] = i;
rusty1s's avatar
rusty1s committed
48
49
50
51
      }
    })
}

rusty1s's avatar
renames  
rusty1s committed
52
53
void index_backward(int dim, THTensor *output, THLongTensor *index, THTensor *grad, THLongTensor *grad_arg) {
  TH_TENSOR_DIM_APPLY4(real, output, int64_t, index, real, grad, int64_t, grad_arg, dim,
rusty1s's avatar
rusty1s committed
54
    for (int64_t i = 0; i < THLongTensor_size(index, dim); i++) {
rusty1s's avatar
renames  
rusty1s committed
55
      if (grad_arg_data[index_data[i]] == i) output_data[i] = grad_data[index_data[i]];
rusty1s's avatar
min max  
rusty1s committed
56
57
58
    })
}

rusty1s's avatar
rusty1s committed
59
#endif