pytorch · guilhermeleobas · May 6, 2025 · May 6, 2025 · May 7, 2025 · May 7, 2025
diff --git a/test/dynamo/test_functions.py b/test/dynamo/test_functions.py
@@ -1747,7 +1747,8 @@ def fn(a):
     def test_set_intersection(a, b):
         set1 = {"apple", "banana", "cherry"}
         set2 = {"google", "microsoft", "apple"}
-        intersection_set = set1.intersection(set2)
+        set3 = {"shoes", "flipflops", "apple"}
+        intersection_set = set1.intersection(set2, set3)
         if "apple" in intersection_set:
             x = a + b
         else:
@@ -1756,7 +1757,31 @@ def test_set_intersection(a, b):
             y = a + b
         else:
             y = a - b
-        return x, y
+        if "shoes" in intersection_set:
+            z = a + b
+        else:
+            z = a - b
+        return x, y, z
+
+    @make_test
+    def test_set_intersection_update(a, b):
+        set1 = {"apple", "banana", "cherry"}
+        set2 = {"google", "microsoft", "apple"}
+        set3 = {"shoes", "flipflops", "apple"}
+        set1.intersection_update(set2, set3)
+        if "apple" in set1:
+            x = a + b
+        else:
+            x = a - b
+        if "banana" in set1:
+            y = a + b
+        else:
+            y = a - b
+        if "shoes" in set1:
+            z = a + b
+        else:
+            z = a - b
+        return x, y, z
 
     @make_test
     def test_set_union(a, b):

diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestOnlySetsDict.test_intersection_update b/test/dynamo_expected_failures/CPython313-test_set-TestOnlySetsDict.test_intersection_update
diff --git a/...dynamo_expected_failures/CPython313-test_set-TestOnlySetsNumeric.test_intersection_update b/...dynamo_expected_failures/CPython313-test_set-TestOnlySetsNumeric.test_intersection_update
diff --git a/test/dynamo_expected_failures/CPython313-test_set-TestOnlySetsTuple.test_intersection_update b/test/dynamo_expected_failures/CPython313-test_set-TestOnlySetsTuple.test_intersection_update
diff --git a/.../dynamo_expected_failures/CPython313-test_set-TestUpdateOps.test_intersection_method_call b/.../dynamo_expected_failures/CPython313-test_set-TestUpdateOps.test_intersection_method_call
diff --git a/torch/_dynamo/polyfills/__init__.py b/torch/_dynamo/polyfills/__init__.py
@@ -125,14 +125,26 @@ def set_isdisjoint(set1, set2):
     return True
 
 
-def set_intersection(set1, set2):
+def set_intersection(set1, *others):
+    if len(others) == 0:
+        return set1.copy()
+
     intersection_set = set()
     for x in set1:
-        if x in set2:
+        for set2 in others:
+            if x not in set2:
+                break
+        else:
             intersection_set.add(x)
     return intersection_set
 
 
+def set_intersection_update(set1, *others):
+    result = set1.intersection(*others)
+    set1.clear()
+    set1.update(result)
+
+
 def set_union(set1, set2):
     union_set = set1.copy()
     set_update(union_set, set2)

diff --git a/torch/_dynamo/variables/dicts.py b/torch/_dynamo/variables/dicts.py
@@ -776,10 +776,14 @@ def call_method(
             ).call_function(tx, [self, args[0]], {})
         elif name == "intersection":
             assert not kwargs
-            assert len(args) == 1
             return variables.UserFunctionVariable(
                 polyfills.set_intersection
-            ).call_function(tx, [self, args[0]], {})
+            ).call_function(tx, [self, *args], {})
+        elif name == "intersection_update":
+            assert not kwargs
+            return variables.UserFunctionVariable(
+                polyfills.set_intersection_update
+            ).call_function(tx, [self, *args], {})
         elif name == "union":
             assert not kwargs
             assert len(args) == 1