Remove unnecessary Subtensor dependencies in nonzero functions

tests/tensor/test_basic.py

@@ -3019,97 +3019,87 @@ class TestTriangle:
 
 
 class TestNonzero:
+    @change_flags(compute_test_value="raise")
     def test_nonzero(self):
         def check(m):
             m_symb = theano.tensor.tensor(
                 dtype=m.dtype, broadcastable=(False,) * m.ndim
             )
+            m_symb.tag.test_value = m
 
-            f_tuple = function([m_symb], nonzero(m_symb, return_matrix=False))
-            f_matrix = function([m_symb], nonzero(m_symb, return_matrix=True))
+            res_tuple_tt = nonzero(m_symb, return_matrix=False)
+            res_matrix_tt = nonzero(m_symb, return_matrix=True)
 
-            assert np.allclose(f_matrix(m), np.vstack(np.nonzero(m)))
-            for i, j in zip(f_tuple(m), np.nonzero(m)):
+            res_tuple = tuple(r.tag.test_value for r in res_tuple_tt)
+            res_matrix = res_matrix_tt.tag.test_value
+
+            assert np.allclose(res_matrix, np.vstack(np.nonzero(m)))
+
+            for i, j in zip(res_tuple, np.nonzero(m)):
                 assert np.allclose(i, j)
 
-        rand0d = np.array(rand())
+        rand0d = np.empty(())
         with pytest.raises(ValueError):
             check(rand0d)
 
-        rand1d = rand(8)
+        rand1d = np.empty((8,))
         rand1d[:4] = 0
         check(rand1d)
 
-        rand2d = rand(8, 9)
+        rand2d = np.empty((8, 9))
         rand2d[:4] = 0
         check(rand2d)
 
-        rand3d = rand(8, 9, 10)
-        rand3d[:4] = 0
-        check(rand3d)
-
-        rand4d = rand(8, 9, 10, 11)
-        rand4d[:4] = 0
-        check(rand4d)
-
+    @change_flags(compute_test_value="raise")
     def test_flatnonzero(self):
         def check(m):
             m_symb = theano.tensor.tensor(
                 dtype=m.dtype, broadcastable=(False,) * m.ndim
             )
-            f = function([m_symb], flatnonzero(m_symb))
-            result = f(m)
+            m_symb.tag.test_value = m
+
+            res_tt = flatnonzero(m_symb)
+
+            result = res_tt.tag.test_value
             assert np.allclose(result, np.flatnonzero(m))
 
-        rand0d = np.array(rand())
+        rand0d = np.empty(())
         with pytest.raises(ValueError):
             check(rand0d)
 
-        rand1d = rand(8)
+        rand1d = np.empty((8,))
         rand1d[:4] = 0
         check(rand1d)
 
-        rand2d = rand(8, 9)
+        rand2d = np.empty((8, 9))
         rand2d[:4] = 0
         check(rand2d)
 
-        rand3d = rand(8, 9, 10)
-        rand3d[:4] = 0
-        check(rand3d)
-
-        rand4d = rand(8, 9, 10, 11)
-        rand4d[:4] = 0
-        check(rand4d)
-
+    @change_flags(compute_test_value="raise")
     def test_nonzero_values(self):
         def check(m):
             m_symb = theano.tensor.tensor(
                 dtype=m.dtype, broadcastable=(False,) * m.ndim
             )
-            f = function([m_symb], nonzero_values(m_symb))
-            result = f(m)
+            m_symb.tag.test_value = m
+
+            res_tt = nonzero_values(m_symb)
+
+            result = res_tt.tag.test_value
             assert np.allclose(result, m[np.nonzero(m)])
 
-        rand0d = rand()
+        rand0d = np.empty(())
         with pytest.raises(ValueError):
             check(rand0d)
 
-        rand1d = rand(8)
+        rand1d = np.empty((8,))
         rand1d[:4] = 0
         check(rand1d)
 
-        rand2d = rand(8, 9)
+        rand2d = np.empty((8, 9))
         rand2d[:4] = 0
         check(rand2d)
 
-        rand3d = rand(8, 9, 10)
-        rand3d[:4] = 0
-        check(rand3d)
-
-        rand4d = rand(8, 9, 10, 11)
-        rand4d[:4] = 0
-        check(rand4d)
-
 
 def test_identity():
     def check(dtype):

theano/tensor/basic.py

@@ -2723,22 +2723,15 @@ class Nonzero(gof.Op):
     """
     Return the indices of the elements that are non-zero.
 
-    Returns a matrix of shape (ndim, number of nonzero elements) such that
-    element (i,j) is the index in the ith dimension of the jth non-zero
-    element.
-
-    Note this is different than NumPy, which returns a tuple of arrays, one for
-    each dimension of the input array.
-
     Parameters
     ----------
-    a : array_like
+    a: array_like
         Input array.
 
     Returns
     -------
-    matrix
-        Matrix containing the indices of the non-zero elements of a.
+    indices: list
+        A list containing the indices of the non-zero elements of `a`.
 
     See Also
     --------
@@ -2754,20 +2747,17 @@ class Nonzero(gof.Op):
         a = as_tensor_variable(a)
         if a.ndim == 0:
             raise ValueError("Nonzero only supports non-scalar arrays.")
-        output = [TensorType(dtype="int64", broadcastable=(False, False))()]
+        output = [
+            TensorType(dtype="int64", broadcastable=(False,))() for i in range(a.ndim)
+        ]
         return gof.Apply(self, [a], output)
 
     def perform(self, node, inp, out_):
         a = inp[0]
-        (out,) = out_
 
         result_tuple = np.nonzero(a)
-        if len(result_tuple[0]) > 0:
-            result = np.vstack(result_tuple)
-        else:
-            result = np.zeros((len(result_tuple), 0))
-
-        out[0] = result.astype("int64")
+        for i, res in enumerate(result_tuple):
+            out_[i][0] = res.astype("int64")
 
     def grad(self, inp, grads):
         return [grad_undefined(self, 0, inp[0])]
@@ -2809,22 +2799,23 @@ def nonzero(a, return_matrix=False):
         flattened input array.
 
     """
-    matrix_result = _nonzero(a)
+    res = _nonzero(a)
+    if isinstance(res, list):
+        res = tuple(res)
+    else:
+        res = (res,)
+
     if return_matrix:
-        return matrix_result
+        if len(res) > 1:
+            return stack(res, 0)
+        elif len(res) == 1:
+            return shape_padleft(res[0])
     else:
-        if a.ndim > 0:
-            tuple_result = tuple([matrix_result[i] for i in range(a.ndim)])
-        else:
-            tuple_result = tuple([matrix_result[0]])
-        return tuple_result
+        return res
 
 
 def flatnonzero(a):
-    """
-    Return a vector of indices that are non-zero in the flattened version of a.
-
-    This is equivalent to nonzero(a.flatten(), return_matrix=True)[0]
+    """Return a vector of indices that are non-zero in the flattened version of `a`.
 
     Parameters
     ----------
@@ -2845,26 +2836,11 @@ def flatnonzero(a):
     """
     if a.ndim == 0:
         raise ValueError("Nonzero only supports non-scalar arrays.")
-    return nonzero(a.flatten(), return_matrix=True)[0]
+    return nonzero(a.flatten(), return_matrix=False)[0]
 
 
 def nonzero_values(a):
-    """
-    Return a vector of non-zero elements contained in the input array.
-
-    The following behavior works to extract non-zero elements from an array
-    in NumPy but is *NOT* supported by Theano:
-
-        a[numpy.nonzero(a)]
-
-    Instead, the nonzero_values function or method should be used:
-
-        tensor.nonzero_values(a)
-        a.nonzero_values()
-
-    This is equivalent to the following:
-
-        a.flatten()[tensor.flatnonzero(a)]
+    """Return a vector of non-zero elements contained in the input array.
 
     Parameters
     ----------