grad basis cpu check

aten/TH/THBasis.c

@@ -24,7 +24,7 @@
         wiOffset *= kernelSizeData[d]; \
 \
         v -= floor(v); \
-        CODE \
+        v = CODE; \
         b *= v; \
       } \
       basisData[e * basis->stride[0] + s * basis->stride[1]] = b; \
@@ -33,5 +33,57 @@
   } \
 }
 
+#define TH_TENSOR_BASIS_BACKWARD(M, self, gradBasis, pseudo, kernelSize, isOpenSpline, CODE, \
+                                 GRAD_CODE) { \
+  real *selfData = THTensor_(data)(self); \
+  real *gradBasisData = THTensor_(data)(gradBasis); \
+  real *pseudoData = THTensor_(data)(pseudo); \
+  int64_t *kernelSizeData = THLongTensor_data(kernelSize); \
+  uint8_t *isOpenSplineData = THByteTensor_data(isOpenSpline); \
+\
+  ptrdiff_t e, d, s; \
+  int64_t k, kMod, wi, wiOffset; \
+  real b, v; \
+  real g; \
+\
+  for (e = 0; e < THTensor_(size)(pseudo, 0); e++) { \
+    for (d = 0; d < THTensor_(size)(pseudo, 1); d++) { \
+      real g_out = 0; \
+      int64_t quotient = pow(M + 1, d); \
+      for (s = 0; s < THTensor_(size)(gradBasis, 1); s++) { \
+        kMod = (s / quotient) % (M + 1); \
+        v = pseudoData[e * pseudo->stride[0] + d * pseudo->stride[1]]; \
+        v *= kernelSizeData[d] - M * isOpenSplineData[d]; \
+        v -= floor(v); \
+        v = GRAD_CODE; \
+        g = v; \
+\
+        ptrdiff_t d_it; \
+        for (d_it = 0; d_it < d; d_it++) { \
+          int64_t quotient2 = pow(M + 1, d_it); \
+          kMod = (s / quotient2) % (M + 1); \
+          v = pseudoData[e * pseudo->stride[0] + d_it * pseudo->stride[1]]; \
+          v *= kernelSizeData[d_it] - M * isOpenSplineData[d_it]; \
+          v -= floor(v); \
+          v = CODE; \
+          g *= v; \
+        } \
+        for (d_it = d + 1; d_it < THTensor_(size)(pseudo, 1); d_it++) { \
+          int64_t quotient2 = pow(M + 1, d_it); \
+          kMod = (s / quotient2) % (M + 1); \
+          v = pseudoData[e * pseudo->stride[0] + d_it * pseudo->stride[1]]; \
+          v *= kernelSizeData[d_it] - M * isOpenSplineData[d_it]; \
+          v -= floor(v); \
+          v = CODE; \
+          g *= v; \
+        } \
+        g_out += g * gradBasisData[e * gradBasis->stride[0] + s * gradBasis->stride[1]]; \
+      } \
+      g_out *= kernelSizeData[d] - M * isOpenSplineData[d]; \
+      selfData[e * self->stride[0] + d * self->stride[1]] = g_out; \
+    } \
+  } \
+}
+
 #include "generic/THBasis.c"
 #include "THGenerateFloatTypes.h"

aten/TH/generic/THBasis.c

@@ -6,12 +6,22 @@ inline real THTensor_(linear)(real v, int64_t kMod) {
   return 1 - v - kMod + 2 * v * kMod;
 }
 
+inline real THTensor_(gradLinear)(real v, int64_t kMod) {
+  return 2 * kMod - 1;
+}
+
 inline real THTensor_(quadratic)(real v, int64_t kMod) {
   if (kMod == 0) return 0.5 * v * v - v + 0.5;
   else if (kMod == 1) return -v * v + v + 0.5;
   else return 0.5 * v * v;
 }
 
+inline real THTensor_(gradQuadratic)(real v, int64_t kMod) {
+  if (kMod == 0) return v - 1;
+  else if (kMod == 1) return -2 * v + 1;
+  else return v;
+}
+
 inline real THTensor_(cubic)(real v, int64_t kMod) {
   if (kMod == 0) { v = (1 - v); return v * v * v / 6.0; }
   else if (kMod == 1) return (3 * v * v * v - 6 * v * v + 4) / 6;
@@ -19,69 +29,47 @@ inline real THTensor_(cubic)(real v, int64_t kMod) {
   else return v * v * v / 6;
 }
 
+inline real THTensor_(gradCubic)(real v, int64_t kMod) {
+  if (kMod == 0) return (-v * v + 2 * v - 1) / 2;
+  else if (kMod == 1) return (3 * v * v -4 * v) / 2;
+  else if (kMod == 2) return (-3 * v * v + 2 * v + 1) / 2;
+  else return v * v / 2;
+}
+
 void THTensor_(linearBasisForward)(THTensor *basis, THLongTensor *weightIndex, THTensor *pseudo,
                                    THLongTensor *kernelSize, THByteTensor *isOpenSpline) {
   TH_TENSOR_BASIS_FORWARD(1, basis, weightIndex, pseudo, kernelSize, isOpenSpline,
-    v = THTensor_(linear)(v, kMod);
-  )
+                          THTensor_(linear)(v, kMod))
 }
 
 void THTensor_(quadraticBasisForward)(THTensor *basis, THLongTensor *weightIndex, THTensor *pseudo,
                                       THLongTensor *kernelSize, THByteTensor *isOpenSpline) {
   TH_TENSOR_BASIS_FORWARD(2, basis, weightIndex, pseudo, kernelSize, isOpenSpline,
-    v = THTensor_(quadratic)(v, kMod);
-  )
+                          THTensor_(quadratic)(v, kMod))
 }
 
 void THTensor_(cubicBasisForward)(THTensor *basis, THLongTensor *weightIndex, THTensor *pseudo,
                                   THLongTensor *kernelSize, THByteTensor *isOpenSpline) {
   TH_TENSOR_BASIS_FORWARD(3, basis, weightIndex, pseudo, kernelSize, isOpenSpline,
-    v = THTensor_(cubic)(v, kMod);
-  )
+                          THTensor_(cubic)(v, kMod))
 }
 
 void THTensor_(linearBasisBackward)(THTensor *self, THTensor *gradBasis, THTensor *pseudo,
                                     THLongTensor *kernelSize, THByteTensor *isOpenSpline) {
-  THTensor_(fill)(self, 0);
-
-  real *selfData = THTensor_(data)(self);
-  real *gradBasisData = THTensor_(data)(gradBasis);
-  real *pseudoData = THTensor_(data)(pseudo);
-  int64_t *kernelSizeData = THLongTensor_data(kernelSize);
-  uint8_t *isOpenSplineData = THByteTensor_data(isOpenSpline);
-
-  ptrdiff_t e, d, s;
-  int64_t k, kMod, wi, wiOffset;
-  real b, v;
-  real g;
-
-  for (e = 0; e < THTensor_(size)(pseudo, 0); e++) {
-    for (d = 0; d < THTensor_(size)(pseudo, 1); d++) {
-      int64_t quotient = pow(2, d);
-      /* printf("e = %i, d = %i, stride0 = %i, stride1 = %i \n", e, d, pseudo->stride[0], pseudo->stride[1]); */
-      for (s = 0; s < THTensor_(size)(gradBasis, 1); s++) {
-        kMod = (s / quotient) % 2;
-        real v = pseudoData[e * pseudo->stride[0] + d * pseudo->stride[1]];
-        v *= kernelSizeData[d] - 2 * isOpenSplineData[d];
-        v -= floor(v);
-        v = -1 + kMod + kMod;  // grad code
-        g = v;
-
-        ptrdiff_t d_it;
-        for (d_it = 0; d_it < d; d_it++) {
-        }
-      }
-    }
-    /* selfData[e * self->stride[0] + d * self->stride[1]] = 1; */
-  }
+  TH_TENSOR_BASIS_BACKWARD(1, self, gradBasis, pseudo, kernelSize, isOpenSpline,
+                           THTensor_(linear)(v, kMod), THTensor_(gradLinear)(v, kMod))
 }
 
 void THTensor_(quadraticBasisBackward)(THTensor *self, THTensor *gradBasis, THTensor *pseudo,
                                        THLongTensor *kernelSize, THByteTensor *isOpenSpline) {
+  TH_TENSOR_BASIS_BACKWARD(2, self, gradBasis, pseudo, kernelSize, isOpenSpline,
+                           THTensor_(quadratic)(v, kMod), THTensor_(gradQuadratic)(v, kMod))
 }
 
 void THTensor_(cubicBasisBackward)(THTensor *self, THTensor *gradBasis, THTensor *pseudo,
                                    THLongTensor *kernelSize, THByteTensor *isOpenSpline) {
+  TH_TENSOR_BASIS_BACKWARD(3, self, gradBasis, pseudo, kernelSize, isOpenSpline,
+                           THTensor_(cubic)(v, kMod), THTensor_(gradCubic)(v, kMod))
 }
 
 #endif // TH_GENERIC_FILE

test/test_basis.py

@@ -57,12 +57,11 @@ def test_spline_basis_gpu(tensor, i):  # pragma: no cover
 
 
 def test_spline_basis_grad_cpu():
-    degree = 1
     kernel_size = torch.LongTensor([5, 5, 5])
     is_open_spline = torch.ByteTensor([1, 0, 1])
-    op = SplineBasis(degree, kernel_size, is_open_spline)
-
     pseudo = torch.DoubleTensor(4, 3).uniform_(0, 1)
     pseudo = Variable(pseudo, requires_grad=True)
 
+    for degree in implemented_degrees.keys():
+        op = SplineBasis(degree, kernel_size, is_open_spline)
         assert gradcheck(op, (pseudo, ), eps=1e-6, atol=1e-4) is True