ENH: array_namespace: support torch.compile (#413)

* ENH: array_namespace: support `torch.compile` without graph breaks

torch.dynamo does not know that module objects are hashable, and graph breaks on compiling a set of module variables. Use lists instead to avoid a graph break.

---------

Co-authored-by: Evgeni Burovski <evgeny.burovskiy@gmail.com>

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,7 +663,19 @@ 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(ns)
+            names.remove(name)
+    namespaces = unique_namespaces
 
     try:
         (xp,) = 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):

tests/test_torch.py

@@ -161,3 +161,18 @@ def test_round():
     r = xp.round(x, decimals=1, out=o)
     assert xp.all(r == o)
     assert r is o
+
+
+def test_dynamo_array_namespace():
+    """Check that torch.compiling array_namespace does not incur graph breaks."""
+    from array_api_compat import array_namespace
+
+    def foo(x):
+        xp = array_namespace(x)
+        return xp.multiply(x, x)
+
+    bar = torch.compile(fullgraph=True)(foo)
+
+    x = torch.arange(3)
+    y = bar(x)
+    assert xp.all(y == x**2)