Introducing array-like sequence methods __contains__ (#17733)

bhushan23 · facebook-github-bot · commit b57fe3cc66cf · 2019-03-11T09:00:16.000-07:00
Summary: for tensor Fixes: #17000 Pull Request resolved: #17733 Differential Revision: D14401952 Pulled By: soumith fbshipit-source-id: c841b128c5a1fceda1094323ed4ef1d0cf494909
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -8151,6 +8151,17 @@ def test_reversed(self):
         val = torch.tensor(42)
         self.assertEqual(reversed(val), torch.tensor(42))
 
+    def test_contains(self):
+        x = torch.arange(0, 10)
+        self.assertEqual(4 in x, True)
+        self.assertEqual(12 in x, False)
+
+        x = torch.arange(1, 10).view(3, 3)
+        val = torch.arange(1, 4)
+        self.assertEqual(val in x, True)
+        val += 10
+        self.assertEqual(val in x, False)
+
     @staticmethod
     def _test_rot90(self, use_cuda=False):
         device = torch.device("cuda" if use_cuda else "cpu")
diff --git a/torch/tensor.py b/torch/tensor.py
@@ -7,6 +7,7 @@
 import weakref
 from torch._six import imap
 from torch._C import _add_docstr
+from numbers import Number
 
 
 # NB: If you subclass Tensor, and want to share the subclassed class
@@ -426,6 +427,17 @@ def __array_wrap__(self, array):
             array = array.astype('uint8')
         return torch.from_numpy(array)
 
+    def __contains__(self, element):
+        r"""Check if `element` is present in tensor
+
+        Arguments:
+            element (Tensor or scalar): element to be checked
+                for presence in current tensor"
+        """
+        if isinstance(element, (torch.Tensor, Number)):
+            return (element == self).any().item()
+        return NotImplemented
+
     @property
     def __cuda_array_interface__(self):
         """Array view description for cuda tensors.
