Merge pull request #30 from deepmind/fix_erf

Fix problem in calling row vectors

luasrc/error_handling.lua

@@ -264,7 +264,7 @@ end
 
 Process the optional return storage, the sizes of the parameter functions, etc
 
-@param K number of actual parameters for the sampler
+@param K number of actual parameters required by the sampler
 @param defaultResultType Tensor class corresponding to the expected result type (e.g. torch.DoubleTensor, torch.IntegerTensor, etc)
 @param ... List of all parameters passed to the original caller
 
@@ -339,8 +339,10 @@ function cephes._check1DParams(K, defaultResultType, ...)
         result:resize(Nresult)
     end
 
+    -- Expand parameters which are of the wrong size. Note: they have
+    -- to be single-element to be expanded
     for paramIndex, param in ipairs(params) do
-        if param:size(1) == 1 then
+        if param:nElement() == 1 then
             local sizes = param:size()
             sizes[1] = Nparams
             params[paramIndex] = params[paramIndex]:expand(sizes)

luasrc/tests/test_vectorized.lua

@@ -125,5 +125,50 @@ function vectorizeTests.testNdtr()
     tester:asserteq(result:size(1), n, "should get " .. n .. " results")
 end
 
+
+function vectorizeTests.testSingleRowInput()
+  local x = torch.Tensor({{1, 2}})
+  local a = cephes.erf(x)
+  tester:assertTensorEq(a,
+                        torch.Tensor{cephes.erf(1),
+                        cephes.erf(2)},
+                        1e-16,
+                        'Wrong output')
+end
+
+
+function vectorizeTests.testSingleRowInputWithResult()
+  local x = torch.Tensor({{1, 2}})
+  local a = torch.Tensor(2)
+  cephes.erf(a, x)
+  tester:assertTensorEq(a,
+                        torch.Tensor{cephes.erf(1),
+                        cephes.erf(2)},
+                        1e-16,
+                        'Wrong output')
+end
+
+
+function vectorizeTests.testSingleColumnInput()
+  local x = torch.Tensor{{1}, {2}}
+  local a = cephes.erf(x)
+  tester:assertTensorEq(a,
+                        torch.Tensor{cephes.erf(1),
+                        cephes.erf(2)},
+                        1e-16,
+                        'Wrong output')
+end
+
+
+function vectorizeTests.testSingleColumnInputWithResult()
+  local x = torch.Tensor{{1}, {2}}
+  local a = torch.Tensor(2)
+  cephes.erf(a, x)
+  tester:assertTensorEq(a,
+                        torch.Tensor{cephes.erf(1),
+                        cephes.erf(2)},
+                        1e-16,
+                        'Wrong output')
+end
 tester:add(vectorizeTests)
 return tester:run()