[Dynamo] Add itertools.compress() support (pytorch#139061)

Use polyfill to add `itertools.compress()` support in Dynamo.

Pull Request resolved: pytorch#139061
Approved by: https://github.com/jansel

Author: shink

test/dynamo/test_functions.py

@@ -326,6 +326,22 @@ def test_itertools_pairwise(a):
             pairs.append(torch.ones(size))
         return pairs
 
+    def test_itertools_compress(self):
+        def fn():
+            return itertools.compress("ABCDEF", [1, 0, 1, 0, 1, 1])
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertListEqual(list(opt_fn()), list(fn()))
+
+    def test_itertools_compress_tensors(self):
+        def fn():
+            return itertools.compress(
+                [torch.tensor([0]), torch.tensor([1]), torch.tensor([2])], [1, 0, 1]
+            )
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertListEqual(list(opt_fn()), list(fn()))
+
     @make_test
     def test_np_iinfo(a):
         max_dim = np.iinfo(np.int16).max

torch/_dynamo/polyfills/itertools.py

@@ -6,7 +6,7 @@
 
 import itertools
 import sys
-from typing import Iterable, Iterator, TypeVar
+from typing import Generator, Iterable, Iterator, TypeVar
 
 from ..decorators import substitute_in_graph
 
@@ -16,10 +16,12 @@
     "chain_from_iterable",
     "islice",
     "tee",
+    "compress",
 ]
 
 
 _T = TypeVar("_T")
+_U = TypeVar("_U")
 
 
 # Reference: https://docs.python.org/3/library/itertools.html#itertools.chain
@@ -101,3 +103,11 @@ def _tee(link) -> Iterator[_T]:  # type: ignore[no-untyped-def]
             return
 
     return tuple(_tee(shared_link) for _ in range(n))
+
+
+# Reference: https://docs.python.org/3/library/itertools.html#itertools.compress
+@substitute_in_graph(itertools.compress, is_embedded_type=True)  # type: ignore[arg-type]
+def compress(
+    data: Iterable[_T], selectors: Iterable[_U], /
+) -> Generator[_T, None, None]:
+    return (datum for datum, selector in zip(data, selectors) if selector)