`array_namespace` incompatible with `torch.compile` due to using modules as dictionary key

[lucascolley]

E       torch._dynamo.exc.UncapturedHigherOrderOpError: while_loop doesn't work unless it is captured completely with torch.compile. Got Dynamo cannot determine whether the underlying object is hashable
E         Explanation: Dynamo does not know whether the underlying python object for PythonModuleVariable(<module 'scipy._external.array_api_compat.torch' from '/Users/lucascolley/ghq/github.com/scipy/scipy/build-cpu-install/usr/lib/python3.13/site-packages/scipy/_external/array_api_compat/torch/__init__.py'>) is hashable
E         Hint: Consider using a different type of object as the dictionary key instead of <class 'module'>.
E         Hint: It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues.
E
E         Developer debug context: is_python_hashable PythonModuleVariable(<module 'scipy._external.array_api_compat.torch' from '/Users/lucascolley/ghq/github.com/scipy/scipy/build-cpu-install/usr/lib/python3.13/site-packages/scipy/_external/array_api_compat/torch/__init__.py'>)
E
E        For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0365.html
E
E       from user code:
E          File "/Users/lucascolley/ghq/github.com/scipy/scipy/.pixi/envs/torch-cpu/lib/python3.13/site-packages/torch/_higher_order_ops/while_loop.py", line 244, in _while_loop_op_wrapper
E           return while_loop_op(*args, **kwargs)
E         File "/Users/lucascolley/ghq/github.com/scipy/scipy/.pixi/envs/torch-cpu/lib/python3.13/site-packages/torch/utils/_device.py", line 109, in __torch_function__
E           return func(*args, **kwargs)
E         File "/Users/lucascolley/ghq/github.com/scipy/scipy/.pixi/envs/torch-cpu/lib/python3.13/site-packages/torch/_higher_order_ops/while_loop.py", line 218, in flat_body_fn
E           return body_fn(*carried, *additional)
E         File "/Users/lucascolley/ghq/github.com/scipy/scipy/build-cpu-install/usr/lib/python3.13/site-packages/scipy/sparse/linalg/_isolve/iterative.py", line 436, in body_fun
E           converged = xp_vector_norm(r, axis=-1) <= atol
E         File "/Users/lucascolley/ghq/github.com/scipy/scipy/build-cpu-install/usr/lib/python3.13/site-packages/scipy/_lib/_array_api.py", line 458, in xp_vector_norm
E           xp = array_namespace(x) if xp is None else xp
E         File "/Users/lucascolley/ghq/github.com/scipy/scipy/build-cpu-install/usr/lib/python3.13/site-packages/scipy/_lib/_array_api_override.py", line 156, in array_namespace
E           return array_api_compat.array_namespace(*numpy_arrays, *api_arrays)
E         File "/Users/lucascolley/ghq/github.com/scipy/scipy/build-cpu-install/usr/lib/python3.13/site-packages/scipy/_external/array_api_compat/common/_helpers.py", line 666, in array_namespace
E           namespaces.add(xp)

The hint seems to be saying we should avoid using modules as dictionary keys, but it seems that this actually arises from using a set of modules:

array-api-compat/array_api_compat/common/_helpers.py

Lines 644 to 666 in e5e0bab

	namespaces: set[Namespace] = set()
	for x in xs:
	xp, info = _cls_to_namespace(cast(Hashable, type(x)), api_version, use_compat)
	if info is _ClsToXPInfo.SCALAR:
	continue
	if (
	info is _ClsToXPInfo.MAYBE_JAX_ZERO_GRADIENT
	and _is_jax_zero_gradient_array(x)
	):
	xp = _jax_namespace(api_version, use_compat)
	if xp is None:
	get_ns = getattr(x, "__array_namespace__", None)
	if get_ns is None:
	raise TypeError(f"{type(x).__name__} is not a supported array type")
	if use_compat:
	raise ValueError(
	"The given array does not have an array-api-compat wrapper"
	)
	xp = get_ns(api_version=api_version)
	namespaces.add(xp)

[ev-br]

Could you please add a reproducer Lucas? Including how you torch.compile it --- it'd be best to understand if it's incompatible with fullgraph=True or torch.Dynamo at all.

[lucascolley]

Could you please add a reproducer Lucas? Including how you torch.compile it --- it'd be best to understand if it's incompatible with fullgraph=True or torch.Dynamo at all.

I'll open a draft PR to SciPy soon where the full reproducer will be clear. But no, fullgraph=True is not required, this reproduces with:

@xp.compile()
def loop():
    return xp.while_loop(cond_fun, body_fun, initial_value)
x_final, *_ = loop()

where array_namespace is called within body_fun.

(For now this is totally avoidable, I just needed to pass an xp=xp to keep the array_namespace call outside of torch.while_loop. However, I imagine this won't be avoidable when wrapping full public API functions with torch.compile.)

---

EDIT: see scipy/scipy#24956

[ev-br]

Hm. scipy/scipy#24956 is a bit large of a reproducer :-), so let me try something simple:

$ cat dynamo_namespace.py 
import torch
from array_api_compat import array_namespace, torch as xp

def foo(x):
    xp = array_namespace(x)
    return xp.sin(x)

#bar = torch.compile(fullgraph=True)(foo)    # this fails due to a graph break in array_namespace
bar = torch.compile(foo)
y = bar(torch.arange(3))

print(y)

$ python dynamo_namespace.py
tensor([0.0000, 0.8415, 0.9093])

Running with $ TORCH_LOGS="output_code" python dynamo_namespace.py produces some reasonably-looking code, and TORCH_LOGS="+dynamo" gives lots of details about the graph break in array_namespace.

So if you say you're seeing problems even when you allow graph breaks, this must be something about the (experimental?) torch.while_function. Maybe try to distill a smaller MWE (MFE?) and open a torch issue? I can try asking informally first once there's an MWE.

Whether we can make array_namespace not graph break... Don't know TBH. Maybe time-box it and try changing a set into a list or use string names instead of namespaces or something like that. Ultimately though, I won't be too surprised if it's too dynamic.

[lucascolley]

Well scipy/scipy@7aa4d3d is not that big :)

Indeed, torch.while_loop seems to have more strict requirements. The hint is a bit confusing, though — it seems to be mistaking our set for a dict. Maybe the difference is ultimately inconsequential, I don't know. Would be great to have input from someone familiar with Dynamo.

[ev-br]

Applying

$ git diff
diff --git a/array_api_compat/common/_helpers.py b/array_api_compat/common/_helpers.py
index 8a307f9..94e7a9d 100644
--- a/array_api_compat/common/_helpers.py
+++ b/array_api_compat/common/_helpers.py
@@ -641,7 +641,7 @@ def array_namespace(
     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,14 @@ def array_namespace(
                 )
             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)
+    namespaces = [ns for ns in namespaces if ns.__name__ in names]
 
     try:
         (xp,) = namespaces

avoids a graph break in array_namespace:

$ cat dynamo_namespace.py 
import torch
from array_api_compat import array_namespace, torch as xp

def foo(x):
    xp = array_namespace(x)
    return xp.sin(x)

bar = torch.compile(fullgraph=True)(foo)      # <<< no graph breaks
#bar = torch.compile(foo)
y = bar(torch.arange(3))

print(y)

$ python dynamo_namespace.py 
/home/br/miniforge3/envs/array-api/lib/python3.12/site-packages/torch/_dynamo/variables/functions.py:2056: UserWarning: Dynamo detected a call to a `functools.lru_cache`-wrapped function. Dynamo ignores the cache wrapper and directly traces the wrapped function. Silent incorrectness is only a *potential* risk, not something we have observed. Enable TORCH_LOGS="+dynamo" for a DEBUG stack trace.
  torch._dynamo.utils.warn_once(msg)
tensor([0.0000, 0.8415, 0.9093])

I suspect (but did not measure!) that the perf hit from the list/set manipulations is not measurable even for multiple calls with small objects because everything is wrapped into an lru_cache anyways.

[lucascolley]

thanks! Sounds like a good idea to me.