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

Activity

1. ev-br commented on Mar 31, 2026

   

   ev-br

   on Mar 31, 2026

   Member

   More actions

   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.
2. lucascolley commented on Mar 31, 2026

   

   lucascolley

   on Mar 31, 2026

   Last edited by lucascolley

   MemberAuthor

   More actions

   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
3. ev-br commented on Mar 31, 2026

   

   ev-br

   on Mar 31, 2026

   Last edited by ev-br

   Member

   More actions

   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.
4. lucascolley commented on Apr 1, 2026

   

   lucascolley

   on Apr 1, 2026

   MemberAuthor

   More actions

   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.
5. ev-br commented on Apr 1, 2026

   

   ev-br

   on Apr 1, 2026

   Member

   More actions

   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.

   👍React with 👍1Reacted by Lucas Colley
6. lucascolley commented on Apr 1, 2026

   

   lucascolley

   on Apr 1, 2026

   MemberAuthor

   More actions

   thanks! Sounds like a good idea to me.
7. ev-br

   closed this as completedin #413on Apr 3, 2026
8. ev-br

   added this to the 1.15 milestone on Apr 3, 2026
9. ev-br

   added

   enhancementNew feature or requestNew feature or request

   on Apr 3, 2026