add nonzero_values and test, also update nonzero docstring to reflect new indexing behavior

theano/tensor/basic.py

@@ -1818,6 +1818,10 @@ class _tensor_py_operators:
         """See `theano.tensor.nonzero`"""
         return nonzero(self)
 
+    def nonzero_values(self):
+        """See `theano.tensor.nonzero_values`"""
+        return nonzero_values(self)
+
     def sort(self,  axis=-1, kind='quicksort', order=None):
         """See `theano.tensor.sort`"""
         from theano.tensor.sort import sort
@@ -3195,11 +3199,14 @@ class Nonzero(gof.Op):
     Returns a tuple of arrays, one for each dimension of `a`, containing
     the indices of the non-zero elements in that dimension.
 
-    If a is a matrix, the corresponding non-zero values can be obtained with::
+    Note that the following NumPy indexing behavior
+    does *NOT* currently work in Theano:
+
+        a[numpy.nonzero(a)]
 
-        a[nonzero(a)[0], nonzero(a)[1]]
+    Use the nonzero_values function to extract nonzero elements instead:
 
-    Note that this is NOT the same indexing behavior as NumPy.
+        tensor.nonzero_values(a)
 
     Parameters
     ----------
@@ -3211,6 +3218,12 @@ class Nonzero(gof.Op):
     tuple_of_arrays : tuple
         Indices of elements that are non-zero.
 
+    See Also
+    --------
+    nonzero_vaulues : Return the non-zero elements of the input array
+    flatnonzero : Return the indices of the non-zero elements of the
+        flattened input array.
+
     """
     def make_node(self, a):
         a = as_tensor_variable(a)
@@ -3249,9 +3262,46 @@ def flatnonzero(a):
     See Also
     --------
     nonzero : Return the indices of the non-zero elements of the input array.
+    nonzero_vaulues : Return the non-zero elements of the input array
     """
     return nonzero(a.flatten())
 
+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)]
+
+    Parameters
+    ----------
+    a : tensor
+        Input tensor
+
+    Returns
+    -------
+    res : vector
+        Output vector, containing the non-zero elements of a.
+
+    See Also
+    --------
+    nonzero : Return the indices of the non-zero elements of the input array.
+    flatnonzero : Return the indices of the non-zero elements of the
+        flattened input array.
+
+    """
+    return a.flatten()[flatnonzero(a)]
 
 class Tri(gof.Op):
     def __init__(self, dtype=None):

theano/tensor/tests/test_basic.py

@@ -41,7 +41,7 @@ from theano.tensor import (_shared, wvector, bvector, autocast_float_as,
         ScalarFromTensor, TensorFromScalar, dtensor4, Rebroadcast, Alloc,
         dtensor3, SpecifyShape, Mean, IncSubtensor, AdvancedIncSubtensor1,
         itensor3, Tile, AdvancedIncSubtensor, switch, Diagonal, Diag,
-        nonzero, flatnonzero)
+        nonzero, flatnonzero, nonzero_values)
 from theano.tests import unittest_tools as utt
 from theano.printing import debugprint
 
@@ -1950,6 +1950,30 @@ def test_flatnonzero():
     rand4d[rand4d > rand4d.mean()] = 0
     check(rand4d)
 
+def test_nonzero_values():
+    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)
+        assert numpy.allclose(result, m[numpy.nonzero(m)])
+
+    rand1d = rand(8)
+    rand1d[rand1d > rand1d.mean()] = 0
+    check(rand1d)
+
+    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):