Repository: OpenMined/PySyft
Pull request #253 (2017-09-25 18:34:13)

masked_select method and tests added for the TensorBase class
Adds the final feature to issue #60 and closes it


Patch:
diff --git a/syft/tensor.py b/syft/tensor.py
index a5f21740c9..9117e74a13 100644
--- a/syft/tensor.py
+++ b/syft/tensor.py
@@ -1304,6 +1304,12 @@ class TensorBase(object):
         self.data[indices] = value
         return self
 
+    def masked_select(self, mask):
+        """
+        See :func:`tensor.masked_select`
+        """
+        return masked_select(self, mask)
+
     def eq(self, t):
         """Returns a new Tensor having boolean True values where an element of the calling tensor is equal to the second Tensor, False otherwise.
         The second Tensor can be a number or a tensor whose shape is broadcastable with the calling Tensor."""
@@ -1331,3 +1337,21 @@ def mv(tensormat, tensorvector):
             compatible with matrix {} '.format(tensorvector.data.shape, tensormat.data.shape))
     else:
         return TensorBase(np.matmul(tensormat.data, tensorvector.data))
+
+
+def masked_select(tensor, mask):
+    """
+    Returns a new 1D Tensor which indexes the ``input`` Tensor according to the binary mask ``mask``.
+    The shapes of the ``mask`` tensor and the ``input`` tensor don’t need to match, but they must be broadcastable.
+
+    :param tensor: Input tensor
+    :param mask: The binary mask (non-zero is treated as true)
+    :return: 1D output tensor
+    """
+    mask = _ensure_tensorbase(mask)
+    tensor = _ensure_tensorbase(tensor)
+    if tensor.encrypted or mask.encrypted:
+        raise NotImplemented
+    mask_broadcasted, data_broadcasted = np.broadcast_arrays(mask.data, tensor.data)
+    indices = np.where(mask_broadcasted)
+    return TensorBase(data_broadcasted[indices])


Test patch:
diff --git a/tests/test_tensor.py b/tests/test_tensor.py
index 55a16ceab9..fe66a1ff28 100644
--- a/tests/test_tensor.py
+++ b/tests/test_tensor.py
@@ -1020,6 +1020,28 @@ class masked_fill_Tests(unittest.TestCase):
         self.assertTrue(np.array_equal(t, TensorBase([[2, 2, 2], [1, 1, 1]])))
 
 
+class masked_select_Tests(unittest.TestCase):
+    def testMasked_select(self):
+        t = TensorBase(np.arange(0, 6).reshape(2, 3))
+        mask = TensorBase([[0, 0, 0], [1, 1, 0]])
+        self.assertTrue(np.array_equal(tensor.masked_select(t, mask), TensorBase([3, 4])))
+
+    def testMasked_select_broadcasting1(self):
+        t = TensorBase(np.arange(0, 6).reshape(2, 3))
+        mask = TensorBase([[1, 1, 0]])
+        self.assertTrue(np.array_equal(tensor.masked_select(t, mask), TensorBase([0, 1, 3, 4])))
+
+    def testMasked_select_broadcasting2(self):
+        t = TensorBase([2.0])
+        mask = TensorBase([[1, 1, 0]])
+        self.assertTrue(np.array_equal(tensor.masked_select(t, mask), TensorBase([2.0, 2.0])))
+
+    def testTensorBase_Masked_select(self):
+        t = TensorBase(np.arange(0, 6).reshape(2, 3))
+        mask = TensorBase([[1, 1, 0]])
+        self.assertTrue(np.array_equal(t.masked_select(mask), TensorBase([0, 1, 3, 4])))
+
+
 class eqTests(unittest.TestCase):
     def testEqWithTensor(self):
         t1 = TensorBase(np.arange(5))