THCWeighting.cu 1.97 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
6
7
#ifndef THC_GENERIC_FILE
#define THC_GENERIC_FILE "generic/THCWeighting.cu"
#else

void THCTensor_(weightingForward)(THCState *state, THCTensor *self, THCTensor *src,
                                  THCTensor *weight, THCTensor *basis,
                                  THCudaLongTensor *weightIndex) {
rusty1s's avatar
rusty1s committed
8
9
  TH_TENSOR_WEIGHTING(weightingForwardKernel, THCTensor_(nElement)(state, self), self, src, weight,
                      basis, weightIndex)
rusty1s's avatar
rusty1s committed
10
11
12
13
14
}

void THCTensor_(weightingBackwardSrc)(THCState *state, THCTensor *self, THCTensor *gradOutput,
                                      THCTensor *weight, THCTensor *basis,
                                      THCudaLongTensor *weightIndex) {
rusty1s's avatar
rusty1s committed
15
16
  THCTensor *tWeight = THCTensor_(newTranspose)(state, weight, 1, 2);
  weight = THCTensor_(newContiguous)(state, tWeight);
rusty1s's avatar
rusty1s committed
17

rusty1s's avatar
rusty1s committed
18
19
  TH_TENSOR_WEIGHTING(weightingBackwardSrcKernel, THCTensor_(nElement)(state, self), self,
                      gradOutput, weight, basis, weightIndex)
rusty1s's avatar
rusty1s committed
20

rusty1s's avatar
rusty1s committed
21
22
  THCTensor_(free)(state, tWeight);
  THCTensor_(free)(state, weight);
rusty1s's avatar
rusty1s committed
23
24
25
26
27
}

void THCTensor_(weightingBackwardWeight)(THCState *state, THCTensor *self, THCTensor *gradOutput,
                                         THCTensor *src, THCTensor *basis,
                                         THCudaLongTensor *weightIndex) {
rusty1s's avatar
rusty1s committed
28
  THCTensor_(fill)(state, self, ScalarConvert<int, real>::to(0));
rusty1s's avatar
rusty1s committed
29
30
  TH_TENSOR_WEIGHTING(weightingBackwardWeightKernel, THCTensor_(nElement)(state, gradOutput), self,
                      gradOutput, src, basis, weightIndex)
rusty1s's avatar
rusty1s committed
31
32
33
34
35
}

void THCTensor_(weightingBackwardBasis)(THCState *state, THCTensor *self, THCTensor *gradOutput,
                                        THCTensor *src, THCTensor *weight,
                                        THCudaLongTensor *weightIndex) {
rusty1s's avatar
rusty1s committed
36
  THCTensor_(fill)(state, self, ScalarConvert<int, real>::to(0));
rusty1s's avatar
rusty1s committed
37
38
  TH_TENSOR_WEIGHTING(weightingBackwardBasisKernel, THCTensor_(nElement)(state, gradOutput), self,
                      gradOutput, src, weight, weightIndex)
rusty1s's avatar
rusty1s committed
39
40
41
}

#endif // THC_GENERIC_FILE