fix 1D case for Mean/Sum/Max/Min

Author: andresy

Mean.lua

@@ -8,15 +8,23 @@ end
 
 function Mean:updateOutput(input)
    self.output:mean(input, self.dimension)
-   self.output = self.output:select(self.dimension, 1)
+   if self.output:nDimension() > 1 then
+      self.output = self.output:select(self.dimension, 1)
+   end
    return self.output
 end
 
 function Mean:updateGradInput(input, gradOutput)
    local size = gradOutput:size():totable()
    local stride = gradOutput:stride():totable()
-   table.insert(size, self.dimension, input:size(self.dimension))
-   table.insert(stride, self.dimension, 0)
+
+   if input:nDimension() > 1 then
+      table.insert(size, self.dimension, input:size(self.dimension))
+      table.insert(stride, self.dimension, 0)
+   else
+      size[1] = input:size(1)
+      stride[1] = 0
+   end
 
    self.gradInput:resizeAs(gradOutput):copy(gradOutput)
    self.gradInput:mul(1/input:size(self.dimension))

Sum.lua

@@ -11,20 +11,28 @@ function Sum:updateOutput(input)
       self.output = input.new()
    end
    self.output:sum(input, self.dimension)
-   self.output = self.output:select(self.dimension, 1)
+   if self.output:nDimension() > 1 then
+      self.output = self.output:select(self.dimension, 1)
+   end
    return self.output
 end
 
 function Sum:updateGradInput(input, gradOutput)
    local size = gradOutput:size():totable()
    local stride = gradOutput:stride():totable()
-   table.insert(size, self.dimension, input:size(self.dimension))
-   table.insert(stride, self.dimension, 0)
+
+   if input:nDimension() > 1 then
+      table.insert(size, self.dimension, input:size(self.dimension))
+      table.insert(stride, self.dimension, 0)
+   else
+      size[1] = input:size(1)
+      stride[1] = 0
+   end
 
    self.gradInput:set(gradOutput:storage(),
                       1,
                       torch.LongStorage(size),
                       torch.LongStorage(stride))
-                      
+
    return self.gradInput
 end

generic/Max.c

@@ -36,7 +36,8 @@ static int nn_(Max_updateOutput)(lua_State *L)
                        *indices_data = theIndex+1;
                        *output_data = theMax;)
 
-  THTensor_(select)(output, NULL, dimension, 0);
+  if(output->nDimension > 1)
+    THTensor_(select)(output, NULL, dimension, 0);
 
   return 1;
 }
@@ -56,25 +57,32 @@ static int nn_(Max_updateGradInput)(lua_State *L)
   THTensor_(resizeAs)(gradInput, input);
   THTensor_(zero)(gradInput);
 
-  dim = THLongStorage_newWithSize(gradOutput->nDimension+1);
-  str = THLongStorage_newWithSize(gradOutput->nDimension+1);
-  for(i = 0, j =  0; j < gradOutput->nDimension+1; j++)
+  if(input->nDimension > 1)
   {
-    if(j == dimension)
+    dim = THLongStorage_newWithSize(gradOutput->nDimension+1);
+    str = THLongStorage_newWithSize(gradOutput->nDimension+1);
+    for(i = 0, j =  0; j < gradOutput->nDimension+1; j++)
     {
-      dim->data[j] = input->size[dimension];
-      str->data[j] = 0;
-      continue;
+      if(j == dimension)
+      {
+        dim->data[j] = input->size[dimension];
+        str->data[j] = 0;
+        continue;
+      }
+
+      dim->data[j] = gradOutput->size[i];
+      str->data[j] = gradOutput->stride[i];
+      i++;
     }
-
-    dim->data[j] = gradOutput->size[i];
-    str->data[j] = gradOutput->stride[i];
-    i++;
+    gradOutputPlusOneDim = THTensor_(newWithStorage)(gradOutput->storage, gradOutput->storageOffset, dim, str);
+    THLongStorage_free(dim);
+    THLongStorage_free(str);
+  }
+  else
+  {
+    THTensor_(retain)(gradOutput);
+    gradOutputPlusOneDim = gradOutput;
   }
-
-  gradOutputPlusOneDim = THTensor_(newWithStorage)(gradOutput->storage, gradOutput->storageOffset, dim, str);
-  THLongStorage_free(dim);
-  THLongStorage_free(str);
 
   TH_TENSOR_DIM_APPLY3(real, gradInput, real, gradOutputPlusOneDim, real, indices, dimension,
                        gradInput_data[ ((long)(*indices_data)-1)*gradInput_stride ] = *gradOutputPlusOneDim_data;)

generic/Min.c

@@ -36,7 +36,8 @@ static int nn_(Min_updateOutput)(lua_State *L)
                        *indices_data = theIndex+1;
                        *output_data = theMin;)
 
-  THTensor_(select)(output, NULL, dimension, 0);
+  if(output->nDimension > 1)
+    THTensor_(select)(output, NULL, dimension, 0);
 
   return 1;
 }
@@ -56,25 +57,32 @@ static int nn_(Min_updateGradInput)(lua_State *L)
   THTensor_(resizeAs)(gradInput, input);
   THTensor_(zero)(gradInput);
 
-  dim = THLongStorage_newWithSize(gradOutput->nDimension+1);
-  str = THLongStorage_newWithSize(gradOutput->nDimension+1);
-  for(i = 0, j =  0; j < gradOutput->nDimension+1; j++)
+  if(input->nDimension > 1)
   {
-    if(j == dimension)
+    dim = THLongStorage_newWithSize(gradOutput->nDimension+1);
+    str = THLongStorage_newWithSize(gradOutput->nDimension+1);
+    for(i = 0, j =  0; j < gradOutput->nDimension+1; j++)
     {
-      dim->data[j] = input->size[dimension];
-      str->data[j] = 0;
-      continue;
+      if(j == dimension)
+      {
+        dim->data[j] = input->size[dimension];
+        str->data[j] = 0;
+        continue;
+      }
+
+      dim->data[j] = gradOutput->size[i];
+      str->data[j] = gradOutput->stride[i];
+      i++;
     }
-
-    dim->data[j] = gradOutput->size[i];
-    str->data[j] = gradOutput->stride[i];
-    i++;
+    gradOutputPlusOneDim = THTensor_(newWithStorage)(gradOutput->storage, gradOutput->storageOffset, dim, str);
+    THLongStorage_free(dim);
+    THLongStorage_free(str);
+  }
+  else
+  {
+    THTensor_(retain)(gradOutput);
+    gradOutputPlusOneDim = gradOutput;
   }
-
-  gradOutputPlusOneDim = THTensor_(newWithStorage)(gradOutput->storage, gradOutput->storageOffset, dim, str);
-  THLongStorage_free(dim);
-  THLongStorage_free(str);
 
   TH_TENSOR_DIM_APPLY3(real, gradInput, real, gradOutputPlusOneDim, real, indices, dimension,
                        gradInput_data[ ((long)(*indices_data)-1)*gradInput_stride ] = *gradOutputPlusOneDim_data;)

test/test.lua

@@ -595,6 +595,15 @@ end
 -- end
 
 function nntest.Max()
+   -- 1D
+   local ini = math.random(3,7)
+   local input = torch.Tensor(ini):zero()
+   local module = nn.Max(1)
+
+   local err = jac.testJacobian(module,input)
+   mytester:assertlt(err,precision, 'error on state ')
+
+   -- 3D
    local ini = math.random(3,5)
    local inj = math.random(3,5)
    local ink = math.random(3,5)
@@ -610,6 +619,15 @@ function nntest.Max()
 end
 
 function nntest.Min()
+   -- 1D
+   local ini = math.random(3,7)
+   local input = torch.Tensor(ini):zero()
+   local module = nn.Min(1)
+
+   local err = jac.testJacobian(module,input)
+   mytester:assertlt(err,precision, 'error on state ')
+
+   -- 3D
    local ini = math.random(3,5)
    local inj = math.random(3,5)
    local ink = math.random(3,5)
@@ -625,6 +643,15 @@ function nntest.Min()
 end
 
 function nntest.Mean()
+   -- 1D
+   local ini = math.random(3,7)
+   local input = torch.Tensor(ini):zero()
+   local module = nn.Mean(1)
+
+   local err = jac.testJacobian(module,input)
+   mytester:assertlt(err,precision, 'error on state ')
+
+   -- 3D
    local ini = math.random(3,5)
    local inj = math.random(3,5)
    local ink = math.random(3,5)
@@ -1423,6 +1450,15 @@ function nntest.SpatialLPPooling()
 end
 
 function nntest.Sum()
+   -- 1D
+   local ini = math.random(3,7)
+   local input = torch.Tensor(ini):zero()
+   local module = nn.Sum(1)
+
+   local err = jac.testJacobian(module,input)
+   mytester:assertlt(err,precision, 'error on state ')
+
+   -- 3D
    local ini = math.random(3,5)
    local inj = math.random(3,5)
    local ink = math.random(3,5)