Add functional conversion of HardShrink

Author: andreaskoepf

HardShrink.lua

@@ -6,11 +6,20 @@ function HardShrink:__init(lam)
 end
 
 function HardShrink:updateOutput(input)
-   input.nn.HardShrink_updateOutput(self, input)
+   input.THNN.HardShrink_updateOutput(
+      input:cdata(),
+      self.output:cdata(),
+      self.lambda
+   )
    return self.output
 end
 
 function HardShrink:updateGradInput(input, gradOutput)
-   input.nn.HardShrink_updateGradInput(self, input, gradOutput)
+   input.THNN.HardShrink_updateGradInput(
+      input:cdata(),
+      gradOutput:cdata(),
+      self.gradInput:cdata(),
+      self.lambda
+   )
    return self.gradInput
 end

THNN.lua

@@ -55,6 +55,18 @@ TH_API void THNN_(DistKLDivCriterion_updateGradInput)(
           THTensor *target,
           THTensor *gradInput,
           bool sizeAverage);
+
+TH_API void THNN_(HardShrink_updateOutput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *output,
+          real lambda);
+TH_API void THNN_(HardShrink_updateGradInput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *gradOutput,
+          THTensor *gradInput,
+          real lambda);
 ]]
 
 -- THGenerator struct declaration copied from torch7/lib/TH/THRandom.h

init.c

@@ -29,9 +29,6 @@
 #include "generic/Tanh.c"
 #include "THGenerateFloatTypes.h"
 
-#include "generic/HardShrink.c"
-#include "THGenerateFloatTypes.h"
-
 #include "generic/SoftShrink.c"
 #include "THGenerateFloatTypes.h"
 
@@ -157,7 +154,6 @@ int luaopen_libnn(lua_State *L)
   nn_FloatSoftMax_init(L);
   nn_FloatSoftPlus_init(L);
   nn_FloatTanh_init(L);
-  nn_FloatHardShrink_init(L);
   nn_FloatSoftShrink_init(L);
   nn_FloatThreshold_init(L);
   nn_FloatPReLU_init(L);
@@ -202,7 +198,6 @@ int luaopen_libnn(lua_State *L)
   nn_DoubleSoftMax_init(L);
   nn_DoubleSoftPlus_init(L);
   nn_DoubleTanh_init(L);
-  nn_DoubleHardShrink_init(L);
   nn_DoubleSoftShrink_init(L);
   nn_DoubleThreshold_init(L);
   nn_DoublePReLU_init(L);

lib/THNN/generic/HardShrink.c

@@ -2,49 +2,29 @@
 #define TH_GENERIC_FILE "generic/HardShrink.c"
 #else
 
-static int nn_(HardShrink_updateOutput)(lua_State *L)
+void THNN_(HardShrink_updateOutput)(THNNState *state, THTensor *input, THTensor *output, real lambda)
 {
-  THTensor *input = luaT_checkudata(L, 2, torch_Tensor);
-  real lambda = luaT_getfieldchecknumber(L, 1, "lambda");
-  THTensor *output = luaT_getfieldcheckudata(L, 1, "output", torch_Tensor);
-
   THTensor_(resizeAs)(output, input);
-  
-  TH_TENSOR_APPLY2(real, output, real, input,				\
-                   if ((*input_data) > lambda) *output_data = *input_data; \
-                   else if ((*input_data) < -lambda) *output_data = *input_data; \
-                   else *output_data = 0;);
-  return 1;
-}
-
-static int nn_(HardShrink_updateGradInput)(lua_State *L)
-{
-  THTensor *input = luaT_checkudata(L, 2, torch_Tensor);
-  real lambda = luaT_getfieldchecknumber(L, 1, "lambda");
-  THTensor *gradOutput = luaT_checkudata(L, 3, torch_Tensor);
-  THTensor *gradInput = luaT_getfieldcheckudata(L, 1, "gradInput", torch_Tensor);
 
-  THTensor_(resizeAs)(gradInput, input);
-  TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, input,	\
-                   if ((*input_data) > lambda || (*input_data) < -lambda) \
-		     *gradInput_data = (*gradOutput_data);		\
-		   else							\
-		     *gradInput_data = 0;				\
-    );
-  return 1;
+  TH_TENSOR_APPLY2(real, output, real, input,
+    if ((*input_data) > lambda)
+      *output_data = *input_data;
+    else if ((*input_data) < -lambda)
+      *output_data = *input_data;
+    else
+      *output_data = 0;
+  );
 }
 
-static const struct luaL_Reg nn_(HardShrink__) [] = {
-  {"HardShrink_updateOutput", nn_(HardShrink_updateOutput)},
-  {"HardShrink_updateGradInput", nn_(HardShrink_updateGradInput)},
-  {NULL, NULL}
-};
-
-static void nn_(HardShrink_init)(lua_State *L)
+void THNN_(HardShrink_updateGradInput)(THNNState *state, THTensor *input, THTensor *gradOutput, THTensor *gradInput, real lambda)
 {
-  luaT_pushmetatable(L, torch_Tensor);
-  luaT_registeratname(L, nn_(HardShrink__), "nn");
-  lua_pop(L,1);
+  THTensor_(resizeAs)(gradInput, input);
+  TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, input,
+    if ((*input_data) > lambda || (*input_data) < -lambda)
+      *gradInput_data = (*gradOutput_data);
+    else
+      *gradInput_data = 0;
+  );
 }
 
 #endif

lib/THNN/generic/THNN.h

@@ -55,4 +55,16 @@ TH_API void THNN_(DistKLDivCriterion_updateGradInput)(
           THTensor *gradInput,
           bool sizeAverage);
 
+TH_API void THNN_(HardShrink_updateOutput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *output,
+          real lambda);
+TH_API void THNN_(HardShrink_updateGradInput)(
+          THNNState *state,
+          THTensor *input,
+          THTensor *gradOutput,
+          THTensor *gradInput,
+          real lambda);
+
 #endif

lib/THNN/init.c

@@ -15,3 +15,6 @@
 
 #include "generic/DistKLDivCriterion.c"
 #include "THGenerateFloatTypes.h"
+
+#include "generic/HardShrink.c"
+#include "THGenerateFloatTypes.h"