nn.Add works with batches

1 files changed, 38 insertions, 16 deletions

diff --git a/test/test.lua b/test/test.lua
index 022ee8a..0ae0e61 100644
--- a/test/test.lua
+++ b/test/test.lua
@@ -23,28 +23,50 @@ end
 
 
 function nntest.Add()
+   local inj_vals = {math.random(3,5), 1}  -- Also test the inj = 1 spatial case
    local ini = math.random(3,5)
-   local inj = math.random(3,5)
    local ink = math.random(3,5)
-   local input = torch.Tensor(ini,inj,ink):zero()
-   local module = nn.Add(ini*inj*ink)
 
-   local err = jac.testJacobian(module,input)
-   mytester:assertlt(err,precision, 'error on state ')
-   local err = jac.testJacobianParameters(module, input, module.bias, module.gradBias)
-   mytester:assertlt(err,precision, 'error on bias ')
+   for ind, inj in pairs(inj_vals) do
+      local input = torch.Tensor(ini,inj,ink):zero()
+      local module = nn.Add(ini,inj,ink)
 
-   local err = jac.testJacobianUpdateParameters(module, input, module.bias)
-   mytester:assertlt(err,precision, 'error on bias [direct update]')
+      -- 1D
+      local err = jac.testJacobian(module,input)
+      mytester:assertlt(err,precision, 'error on state ')
 
-   for t,err in pairs(jac.testAllUpdate(module, input, 'bias', 'gradBias')) do
-      mytester:assertlt(err, precision, string.format(
-                         'error on bias [%s]', t))
-   end
+      local err = jac.testJacobianParameters(module, input, module.bias, module.gradBias)
+      mytester:assertlt(err,precision, 'error on bias ')
 
-   local ferr,berr = jac.testIO(module,input)
-   mytester:asserteq(ferr, 0, torch.typename(module) .. ' - i/o forward err ')
-   mytester:asserteq(berr, 0, torch.typename(module) .. ' - i/o backward err ')
+      local err = jac.testJacobianUpdateParameters(module, input, module.bias)
+      mytester:assertlt(err,precision, 'error on bias [direct update] ')
+
+      for t,err in pairs(jac.testAllUpdate(module, input, 'bias', 'gradBias')) do
+         mytester:assertlt(err, precision, string.format('error on bias [%s]', t))
+      end
+
+      -- 2D
+      local nframe = math.random(50,70)
+      local input = torch.Tensor(nframe, ini,inj,ink):zero()
+
+      local err = jac.testJacobian(module,input)
+      mytester:assertlt(err,precision, 'error on state ')
+
+      local err = jac.testJacobianParameters(module, input, module.bias, module.gradBias)
+      mytester:assertlt(err,precision, 'error on bias ')
+
+      local err = jac.testJacobianUpdateParameters(module, input, module.bias)
+      mytester:assertlt(err,precision, 'error on bias [direct update] ')
+
+      for t,err in pairs(jac.testAllUpdate(module, input, 'bias', 'gradBias')) do
+         mytester:assertlt(err, precision, string.format('error on bias [%s]', t))
+      end
+
+      -- IO
+      local ferr,berr = jac.testIO(module,input)
+      mytester:asserteq(ferr, 0, torch.typename(module) .. ' - i/o forward err ')
+      mytester:asserteq(berr, 0, torch.typename(module) .. ' - i/o backward err ')
+   end 
 end
 
 function nntest.CMul()