add nonzero test

theano/tensor/tests/test_basic.py

@@ -40,7 +40,8 @@ from theano.tensor import (_shared, wvector, bvector, autocast_float_as,
         tile, patternbroadcast, Eye, Shape, Default, Dot, PermuteRowElements,
         ScalarFromTensor, TensorFromScalar, dtensor4, Rebroadcast, Alloc,
         dtensor3, SpecifyShape, Mean, IncSubtensor, AdvancedIncSubtensor1,
-        itensor3, Tile, AdvancedIncSubtensor, switch, Diagonal, Diag)
+        itensor3, Tile, AdvancedIncSubtensor, switch, Diagonal, Diag,
+        nonzero)
 from theano.tests import unittest_tools as utt
 from theano.printing import debugprint
 
@@ -1895,6 +1896,30 @@ def test_tril_triu():
         yield check_u, m, -1
 
 
+def test_nonzero():
+    def check(m):
+        m_symb = theano.tensor.tensor(dtype=m.dtype,
+                                      broadcastable = (False,) * m.ndim)
+
+        f = function([m_symb], nonzero(m_symb))
+        result = f(m)
+
+        for i, j in zip(result, numpy.nonzero(m)):
+            assert numpy.allclose(i, j)
+
+    rand2d = rand(8, 9)
+    rand2d[rand2d > rand2d.mean()] = 0
+    check(rand2d)
+
+    rand3d = rand(8, 9, 10)
+    rand3d[rand3d > rand3d.mean()] = 0
+    check(rand3d)
+
+    rand4d = rand(8, 9, 10, 11)
+    rand4d[rand4d > rand4d.mean()] = 0
+    check(rand4d)
+
+
 def test_identity():
     def check(dtype):
         obj = rand_of_dtype((2,), dtype)