1 files changed, 51 insertions, 6 deletions

diff --git a/generic/Sqrt.c b/generic/Sqrt.c
index 3e0c3d9..952c260 100644
--- a/generic/Sqrt.c
+++ b/generic/Sqrt.c
@@ -9,10 +9,30 @@ static int nn_(Sqrt_updateOutput)(lua_State *L)
   THTensor *output = luaT_getfieldcheckudata(L, 1, "output", torch_(Tensor_id));
 
   THTensor_(resizeAs)(output, input);
+  
+  if (input->nDimension == 1 || !THTensor_(isContiguous)(input) || !THTensor_(isContiguous)(output))
+  {
+    TH_TENSOR_APPLY2(real, output, real, input,   \
+         *output_data = sqrt(*input_data + bias););
+  }
+  else
+  {
+    real* output_data = THTensor_(data)(output);
+    real* input_data  = THTensor_(data)(input);
+    long k;
 
-  TH_TENSOR_APPLY2(real, output, real, input,		\
-		   *output_data = sqrt(*input_data + bias););
-
+#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] = sqrt(ptr_input[i] + bias);
+      }
+    }
+  }
   return 1;
 }
 
@@ -25,9 +45,34 @@ static int nn_(Sqrt_updateGradInput)(lua_State *L)
 
   THTensor_(resizeAs)(gradInput, input);
 
-  TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, output,	\
-		   *gradInput_data  = 0.5 * (*gradOutput_data / *output_data););
-  
+  if (output->nDimension == 1 || 
+      !THTensor_(isContiguous)(output) || 
+      !THTensor_(isContiguous)(gradOutput) ||
+      !THTensor_(isContiguous)(gradInput))
+  {
+    TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, output, \
+         *gradInput_data  = 0.5 * (*gradOutput_data / *output_data););
+  }
+  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++)
+      {
+        ptr_gradInput[i] = 0.5 * (ptr_gradOutput[i] / ptr_output[i]);
+      }
+    }
+  }
   return 1;
 }