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

5
void spline_(basis_linear)(THTensor *basis, THLongTensor *weight_index, THTensor *pseudo, THLongTensor *kernel_size, THByteTensor *is_open_spline, int K) {
rusty1s's avatar
rusty1s committed
6
7
  int64_t *kernel_size_data = kernel_size->storage->data + kernel_size->storageOffset;
  uint8_t *is_open_spline_data = is_open_spline->storage->data + is_open_spline->storageOffset;
8
9
10
  int64_t D = THTensor_(size)(pseudo, 1);
  int64_t S = THLongTensor_size(weight_index, 1);
  int64_t s, d;
rusty1s's avatar
rusty1s committed
11
12
13
14

  TH_TENSOR_DIM_APPLY3(real, basis, int64_t, weight_index, real, pseudo, 1, TH_TENSOR_DIM_APPLY3_SIZE_EQ_EXCEPT_DIM,
    for (s = 0; s < S; s++) {
      real b = 1; int64_t i = 0;
15
16
      int64_t k = s;
      int64_t bla = K;
rusty1s's avatar
rusty1s committed
17
      for (d = 0; d < D; d++) {
18
19
20
        bla /= kernel_size_data[d];
        int64_t mod = k % 2;
        k >>= 1;
rusty1s's avatar
rusty1s committed
21

22
        real value = *(pseudo_data + d * pseudo_stride) * (kernel_size_data[d] - is_open_spline_data[d]);
rusty1s's avatar
rusty1s committed
23
24
25
26
27
        int64_t bot = ((int64_t) value) % kernel_size_data[d];
        int64_t top = ((int64_t) (value + 1)) % kernel_size_data[d];
        value -= floor(value);

        b *= (1 - mod) * value + mod * (1 - value);
28
        i += ((1 - mod) * top + mod * bot) * bla;
rusty1s's avatar
rusty1s committed
29
30
31
32
      }
      basis_data[s * basis_stride] = b;
      weight_index_data[s * weight_index_stride] = i;
    })
rusty1s's avatar
rusty1s committed
33
34
35
}

#endif