query: `is_*_namespace`: input must be hashable?

[lucascolley]

The lru_cache docstring says:

Arguments to the cached function must be hashable.

But functions like

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

Lines 340 to 341 in 9469947

	@lru_cache(100)
	def is_cupy_namespace(xp: Namespace) -> bool:

take modules as arguments, which are not hashable.

Is this a problem? The is_*_array functions currently cast args to Hashable before calling another lru_cache-wrapped function:

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

Lines 190 to 191 in 9469947

	cls = cast(Hashable, type(x))
	return _issubclass_fast(cls, "ndonnx", "Array")

cc @crusaderky

[ev-br]

Is this a problem?

Modules are hashable, unless under torch.compile IIUC, so I suppose it is related to #411, which you fixed in gh-413.

Here's a quick poking around session with torch.compile below. It looks like Dynamo special-cases @lru_cache so that it compiles at least, but there's some rough edge with resource management (cf the printed backtrace on exit). No idea how severe it is though, so probably some more investigation is in order.

In [1]: from array_api_compat import is_numpy_namespace, is_torch_namespace, is_jax_namespace

In [2]: def foo(x, xp): 
   ...:     x = x*x
   ...:     return x, is_torch_namespace(xp), is_numpy_namespace(xp)
   ...: 

In [4]: import torch

In [5]: bar = torch.compile(fullgraph=True)(foo)

In [6]: bar(torch.arange(3), torch)
Out[6]: (tensor([0, 1, 4]), True, False)

In [7]: exit()
Exception ignored in: <function ExactWeakKeyDictionary.__setitem__.<locals>.<lambda> at 0x7fde1cd06340>
Traceback (most recent call last):
  File "/home/br/miniforge3/envs/scipy-dev/lib/python3.12/site-packages/torch/_dynamo/utils.py", line 908, in <lambda>
    self.refs[idx] = weakref.ref(key, lambda ref: self._remove_id(idx))
                                                  ^^^^^^^^^^^^^^^^^^^^
  File "/home/br/miniforge3/envs/scipy-dev/lib/python3.12/site-packages/torch/_dynamo/utils.py", line 1928, in _remove_id
    hook()
  File "/home/br/miniforge3/envs/scipy-dev/lib/python3.12/site-packages/torch/_dynamo/utils.py", line 1912, in __call__
    del self.scope[self.name]
        ~~~~~~~~~~^^^^^^^^^^^
KeyError: '__builtins_dict___0'

---

Two more comments/questions:

• we could use module names for cache keys, similar to ENH: array_namespace: support torch.compile #413; cf Raymond Hettinger's answer at
  https://stackoverflow.com/questions/30730983/make-lru-cache-ignore-some-of-the-function-arguments which suggests an auxilliary function. Will probably want to see if this
  • helps with the KeyError above;
  • has a measurable perf impact; IIRC @crusaderky added caching to improve perf on some workloads in ENH: cache helper functions #308, what were they?
• not all is_*_namespace functions are decorated with @lru_cache; e.g. is_jax_namespace is not. Was this deliberate?

[lucascolley]

Modules are hashable

ah, okay, this might be a bug in the type-checker I was using, sorry. DetachHead/basedpyright#1779