Merge pull request #2690 from swook/master

theano/tensor/extra_ops.py

@@ -511,6 +511,26 @@ def squeeze(x):
     return view
 
 
+def compress(condition, x, axis=None):
+    """Return selected slices of an array along given axis.
+
+    It returns the input tensor, but with selected slices along a given axis
+    retained. If no axis is provided, the tensor is flattened
+    Corresponds to numpy.compress
+
+    :param x: Input data, tensor variable
+
+    :param condition: 1 dimensional array of non-zero and zero values
+    corresponding to indices of slices along a selected axis
+
+    :return: `x` with selected slices
+
+    .. versionadded:: 0.7
+    """
+    indices = theano.tensor.basic.flatnonzero(condition)
+    return x.take(indices, axis=axis)
+
+
 class RepeatOp(theano.Op):
     # See the repeat function for docstring
 

theano/tensor/tests/test_extra_ops.py

@@ -7,7 +7,7 @@ from theano.tests import unittest_tools as utt
 
 from theano.tensor.extra_ops import (CumsumOp, cumsum, CumprodOp, cumprod,
                                      BinCountOp, bincount, DiffOp, diff,
-                                     squeeze, RepeatOp, repeat,
+                                     squeeze, compress, RepeatOp, repeat,
                                      Bartlett, bartlett,
                                      FillDiagonal, fill_diagonal,
                                      FillDiagonalOffset, fill_diagonal_offset,
@@ -344,6 +344,45 @@ class SqueezeTester(utt.InferShapeTester):
             assert numpy.allclose(tested, expected)
 
 
+class CompressTester(utt.InferShapeTester):
+    axis_list = [None,
+                 -1,
+                 0,
+                 0,
+                 0,
+                 1]
+    cond_list = [[1, 0, 1, 0, 0, 1],
+                 [0, 1, 1, 0],
+                 [0, 1, 1, 0],
+                 [],
+                 [0, 0, 0, 0],
+                 [1, 1, 0, 1, 0]]
+    shape_list = [(2, 3),
+                  (4, 3),
+                  (4, 3),
+                  (4, 3),
+                  (4, 3),
+                  (3, 5)]
+
+    def setUp(self):
+        super(CompressTester, self).setUp()
+        self.op = compress
+
+    def test_op(self):
+        for axis, cond, shape in zip(self.axis_list, self.cond_list, self.shape_list):
+            cond_var = theano.tensor.ivector()
+            data     = numpy.random.random(size=shape).astype(theano.config.floatX)
+            data_var = theano.tensor.matrix()
+
+            f = theano.function([cond_var, data_var], self.op(cond_var, data_var, axis=axis))
+
+            expected = numpy.compress(cond, data, axis=axis)
+            tested = f(cond, data)
+
+            assert tested.shape == expected.shape
+            assert numpy.allclose(tested, expected)
+
+
 class TestRepeatOp(utt.InferShapeTester):
     def _possible_axis(self, ndim):
         return [None] + range(ndim) + [-i for i in range(ndim)]

theano/tensor/var.py

@@ -596,6 +596,11 @@ class _tensor_py_operators:
         """
         return theano.tensor.extra_ops.squeeze(self)
 
+    def compress(self, a, axis=None):
+        """Return selected slices only
+        """
+        return theano.tensor.extra_ops.compress(self, a, axis=axis)
+
 
 class TensorVariable(_tensor_py_operators, Variable):
     """Subclass to add the tensor operators to the basic `Variable` class."""