1 files changed, 52 insertions, 5 deletions

diff --git a/generic/Tanh.c b/generic/Tanh.c
index 5c24d15..01e9bc0 100644
--- a/generic/Tanh.c
+++ b/generic/Tanh.c
@@ -9,9 +9,28 @@ static int nn_(Tanh_updateOutput)(lua_State *L)
 
   THTensor_(resizeAs)(output, input);
 
-  TH_TENSOR_APPLY2(real, output, real, input,   \
-                   *output_data = tanh(*input_data);)
+  if (input->nDimension == 1 || !THTensor_(isContiguous)(input) || !THTensor_(isContiguous)(output))
+  {
+    TH_TENSOR_APPLY2(real, output, real, input,   \
+         *output_data = tanh(*input_data););
+  }
+  else
+  {
+    real* output_data = THTensor_(data)(output);
+    real* input_data  = THTensor_(data)(input);
+    long k;
 
+#pragma omp parallel for private(k)
+    for (k = 0; k < input->size[0]; k++)
+    {
+      real* ptr_output = output_data + k*input->stride[0];
+      real* ptr_input  = input_data  + k*input->stride[0];
+      long i;
+      for (i = 0; i < input->stride[0]; i++)
+  ptr_output[i] = tanh(ptr_input[i]);
+    }
+  }
+      
   return 1;
 }
 
@@ -22,9 +41,37 @@ static int nn_(Tanh_updateGradInput)(lua_State *L)
   THTensor *gradInput = luaT_getfieldcheckudata(L, 1, "gradInput", torch_(Tensor_id));
 
   THTensor_(resizeAs)(gradInput, output);
-  TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, output, \
-                   real z = *output_data;                              \
-                   *gradInput_data = *gradOutput_data * (1. - z*z););
+
+  if (output->nDimension == 1 || 
+      !THTensor_(isContiguous)(output) || 
+      !THTensor_(isContiguous)(gradOutput) ||
+      !THTensor_(isContiguous)(gradInput))
+  {
+    TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, output,  \
+         real z = *output_data;            \
+         *gradInput_data = *gradOutput_data * (1. - z*z););
+  }
+  else
+  {
+    real* gradOutput_data = THTensor_(data)(gradOutput);
+    real* gradInput_data  = THTensor_(data)(gradInput);
+    real* output_data     = THTensor_(data)(output);
+    long k;
+
+#pragma omp parallel for private(k)
+    for (k = 0; k < output->size[0]; k++)
+    {
+      real* ptr_gradOutput = gradOutput_data + k*output->stride[0];
+      real* ptr_gradInput  = gradInput_data  + k*output->stride[0];
+      real* ptr_output     = output_data     + k*output->stride[0];
+      long i;
+      for (i = 0; i < output->stride[0]; i++)
+      {
+  real z = ptr_output[i];
+  ptr_gradInput[i] = ptr_gradOutput[i] * (1. - z*z);
+      }
+    }
+  }
   return 1;
 }