ENH: array_namespace: support torch.compile

Co-authored-by: Evgeni Burovski <[email protected]>

Author: lucascolley

array_api_compat/common/_helpers.py

@@ -641,7 +641,7 @@ def your_function(x, y):
     is_pydata_sparse_array
 
     """
-    namespaces: set[Namespace] = set()
+    namespaces: list[Namespace] = []
     for x in xs:
         xp, info = _cls_to_namespace(cast(Hashable, type(x)), api_version, use_compat)
         if info is _ClsToXPInfo.SCALAR:
@@ -663,10 +663,22 @@ def your_function(x, y):
                 )
             xp = get_ns(api_version=api_version)
 
-        namespaces.add(xp)
+        namespaces.append(xp)
+
+    # Use a list of modules to avoid a graph break under torch.compile:
+    # torch._dynamo.exc.Unsupported: Dynamo cannot determine whether the underlying object is hashable
+    # Explanation: Dynamo does not know whether the underlying python object for
+    # PythonModuleVariable(<module 'array_api_compat.torch' from ...) is hashable
+    names = set(x.__name__ for x in namespaces)
+    unique_namespaces = []
+    for ns in namespaces:
+        if (name := ns.__name__) in names:
+            unique_namespaces.append(namespaces)
+            names.remove(name)
+    namespaces = unique_namespaces
 
     try:
-        (xp,) = namespaces
+        (xp,) = tuple(namespaces)
         return xp
     except ValueError:
         if not namespaces:

tests/test_no_dependencies.py

@@ -17,7 +17,7 @@ class Array:
     # Dummy array namespace that doesn't depend on any array library
     def __array_namespace__(self, api_version=None):
         class Namespace:
-            pass
+            __name__: str = "foobar"
         return Namespace()
 
 def _test_dependency(mod):