From 0375835dc05bd3d0a6e433b57eb20b88be8f70e7 Mon Sep 17 00:00:00 2001 From: Johnnie Gray Date: Sat, 22 Aug 2026 12:11:49 -0700 Subject: [PATCH] composed functions can optionally supply namespace --- autoray/autoray.py | 152 +++++++++++++++++++++++++++++++++---- docs/automatic_dispatch.md | 18 +++++ docs/changelog.md | 2 + tests/test_autoray.py | 132 ++++++++++++++++++++++++++++++++ 4 files changed, 290 insertions(+), 14 deletions(-) diff --git a/autoray/autoray.py b/autoray/autoray.py index 82e5064..c7e5f6d 100644 --- a/autoray/autoray.py +++ b/autoray/autoray.py @@ -982,8 +982,8 @@ def torch_to_numpy(x): _FUNCS.pop((backend, name), None) _FUNCS.pop((backend + "[alt]", name), None) - # remove all namespaces, because each one keeps the functions that it found - _NAMESPACE_CACHE.clear() + # each namespace holds the functions that it found -> drop + _reset_namespaces() # the dtype cache also keeps results from ``get_lib_fn`` _to_backend_dtype_from_str_cached.cache_clear() # this holds arrays built with ``array`` and ``astype`` @@ -1395,17 +1395,42 @@ def tree_apply_dict(f, tree, is_leaf): # --------------------------- composed functions ---------------------------- # +def _choose_namespace(backend, args): + """Choose the namespace to supply to a composed function, given the + already chosen ``backend`` and the positional ``args`` of the call. The + first argument supplies the dtype and device defaults if it belongs to + ``backend``, otherwise the namespace has none. + """ + try: + like = args[0] + except IndexError: + # without an array there is no dtype or device context to inherit + return get_namespace(like=backend) + + # only inherit context from an argument belonging to the chosen backend, + # a string would be read as a backend name rather than an array + if not isinstance(like, str): + if _infer_class_backend_cached(like.__class__) == backend: + return get_namespace(like) + + # backend selection may have been independent of the first argument + return get_namespace(like=backend) + + class Composed: """Compose an ``autoray.do`` using function. See the main wrapper ``compose``. """ + # no __slots__: functools.wraps writes function metadata to each instance def __init__(self, fn, name=None): self._default_fn = fn if name is None: name = fn.__name__ self._name = name - self._supply_backend = "backend" in signature(fn).parameters + parameters = signature(fn).parameters + self._supply_backend = "backend" in parameters + self._supply_namespace = "namespace" in parameters # this registers the fact that when `get_lib_fn` is called, the # function can be created even if it doesn't exist for a specific @@ -1424,16 +1449,48 @@ def wrapper(fn): return wrapper + def _make_default_function(self, backend, namespace=None): + if self._supply_namespace and namespace is None: + # ordinary dispatch knows the backend but not the namespace, so + # infer it per call, keeping array dtype and device defaults + default_fn = self._default_fn + supply_backend = self._supply_backend + + @functools.wraps(default_fn) + def fn(*args, **kwargs): + if "namespace" not in kwargs: + kwargs["namespace"] = _choose_namespace(backend, args) + if supply_backend and "backend" not in kwargs: + kwargs["backend"] = backend + return default_fn(*args, **kwargs) + + # attach the composed function so a namespace lookup can find + # it and rebuild this default with a fixed namespace + fn._autoray_composed = self + return fn + + # collect the special arguments known before the composed body runs + supplied = {} + if self._supply_backend: + supplied["backend"] = backend + if self._supply_namespace: + # a namespace lookup supplies the exact root xp to bind + supplied["namespace"] = namespace + + if not supplied: + return self._default_fn + + # make sure it inherits __name__ etc + return functools.wraps(self._default_fn)( + functools.partial(self._default_fn, **supplied) + ) + def make_function(self, backend): """Make a new function for the specific ``backend``.""" - if self._supply_backend: - # make sure it inherits __name__ etc - fn = functools.wraps(self._default_fn)( - functools.partial(self._default_fn, backend=backend) - ) - else: - fn = self._default_fn - self.register(backend, fn) + fn = self._make_default_function(backend) + # nothing can have looked this up before it existed, so unlike + # ``register`` there is nothing to invalidate + _FUNCS[backend, self._name] = fn return fn def __call__(self, *args, like=None, **kwargs): @@ -1455,7 +1512,11 @@ def compose(fn=None, *, name=None): for specific implementations to be overridden for specific backends. If the function takes a ``backend`` argument, it will be supplied with the - backend name, to save having to re-choose the backend. + backend name, to save having to re-choose the backend. If it takes a + ``namespace`` argument, it will similarly be supplied with an + ``AutoNamespace``. Calling through a namespace supplies that namespace, + otherwise it is taken from the first argument if that matches the + backend, and has no dtype or device defaults if not. Specific implementations can be provided by calling the ``register`` method of the composed function, or it can itself be used like a decorator:: @@ -2550,6 +2611,18 @@ def _get_submodule(self, name): new._submodule = name return new + def _get_root_namespace(self): + # a root namespace is already the object that should be supplied + if self._submodule is None: + return self + + # namespaces are cached, so this is the root this submodule came from + return get_namespace( + like=self._backend, + device=self._device, + dtype=self._dtype, + ) + def _get_fn(self, name): if name.startswith("__") and name.endswith("__"): # raise correct error for dunder methods, so @@ -2568,11 +2641,22 @@ def _get_fn(self, name): return self._get_submodule(name) if self._backend is None: - # use auto dispatch + # defer dispatch, including namespace inference, until the call return DoFunc(name) fn = get_lib_fn(self._backend, name) + # only a default generated by compose has this attribute, a + # registered implementation does not and is called as it is + composed = getattr(fn, "_autoray_composed", None) + if composed is not None: + # rebuild the default with this exact namespace, rather than one + # inferred on every call + fn = composed._make_default_function( + self._backend, + namespace=self._get_root_namespace(), + ) + # possibly wrap for dtype and device injection if name in _CREATION_ROUTINES: key = (self._backend, name) @@ -2618,9 +2702,40 @@ def __repr__(self): ) +# the instance attributes of an ``AutoNamespace``, everything else in its +# ``__dict__`` is a cached function or submodule lookup +_NAMESPACE_ATTRS = ("_backend", "_device", "_dtype", "_submodule") + _NAMESPACE_CACHE = {} +@functools.lru_cache(2**14) +def _namespace_key_part(x): + """Cached ``str`` of a device or dtype, which normalizes them for the + namespace cache key. Cached because ``numpy.dtype.__str__`` is slow, and + a composed function taking a ``namespace`` looks one up on every call. + """ + return str(x) + + +def _reset_namespace(xp): + """Drop the cached lookups of ``xp`` and of any submodule it made.""" + d = xp.__dict__ + xp.__dict__ = {k: d[k] for k in _NAMESPACE_ATTRS} + for v in d.values(): + if isinstance(v, AutoNamespace): + _reset_namespace(v) + + +def _reset_namespaces(): + """Drop the cached function and submodule lookups of every live namespace, + keeping the namespace objects themselves, so that a namespace held by a + caller stays valid when functions are registered. + """ + for xp in _NAMESPACE_CACHE.values(): + _reset_namespace(xp) + + def get_namespace(like=None, device=None, dtype=None, submodule=None): """Get an automatic namespace object. @@ -2648,7 +2763,16 @@ def get_namespace(like=None, device=None, dtype=None, submodule=None): An automatic namespace object. """ backend, device, dtype = infer_backend_device_dtype(like, device, dtype) - key = (backend, str(device), str(dtype), submodule) + try: + key = ( + backend, + _namespace_key_part(device), + _namespace_key_part(dtype), + submodule, + ) + except TypeError: + # unhashable device or dtype + key = (backend, str(device), str(dtype), submodule) try: xp = _NAMESPACE_CACHE[key] except KeyError: diff --git a/docs/automatic_dispatch.md b/docs/automatic_dispatch.md index 6d6c34c..8662eb7 100644 --- a/docs/automatic_dispatch.md +++ b/docs/automatic_dispatch.md @@ -375,6 +375,24 @@ def my_func_numba(x): do("my_func", x_numpy) ``` +If the default implementation has a `namespace` parameter, `compose` supplies +it automatically, which is often faster and tidier than repeated `do` calls: + +```python +@compose +def my_func(x, namespace): + # get how many elements are needed to sum to 20 + return namespace.sum(namespace.cumsum(x, 0) < 20) +``` + +Calling through a namespace, as `xp.my_func(x)`, supplies that namespace, so +its dtype and device defaults apply. Other calls take the namespace from the +first argument if it matches the dispatched backend, and otherwise get one +with no dtype or device defaults. Note that an explicit `like` argument +selects the backend only, it does not supply these defaults. Implementations +registered with `my_func.register(...)` are called with their own signatures +instead. + ### Deviations from `numpy` diff --git a/docs/changelog.md b/docs/changelog.md index cd38837..db7cdd9 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -6,6 +6,8 @@ Release notes for `autoray`. **Enhancements:** +- `compose` now supplies an [`AutoNamespace`](autoray.autoray.AutoNamespace) to any *default* implementation taking a `namespace` parameter: the namespace itself when called as `xp.my_func(...)`, so its dtype and device defaults apply, otherwise one from the first argument if it matches the dispatched backend. Implementations registered for a specific backend keep their own signatures. +- Namespaces from `get_namespace` are now stable objects: registering a function drops their cached lookups rather than discarding the namespaces, so one already held picks up the new function. - `autoray.lazy` `"sum"`, `"prod"`, `"min"` and `"max"` now accept `keepdims`, matching the eager and array API signatures. The reduced axes are kept as size 1 in the inferred lazy shape. - Added the lazy reductions `"mean"`, `"std"`, `"var"`, `"all"`, `"any"`, `"count_nonzero"`, `"argmin"` and `"argmax"`, which also take `axis` and `keepdims`, and `"cumsum"`, which takes `axis` and accumulates over the flattened array when it is `None`. `"argmin"` and `"argmax"` follow the eager convention of only accepting a scalar `axis`. Backend support for `keepdims` varies: see `XFAILS` in `tests/conftest.py`. - The lazy reductions and `"cumsum"` pass any further keyword arguments, such as `ddof` or `dtype`, straight through to the backend function. A `LazyArray` supplied this way, for example as `where`, is tracked as a dependency of the result. diff --git a/tests/test_autoray.py b/tests/test_autoray.py index b8f88d7..680963f 100644 --- a/tests/test_autoray.py +++ b/tests/test_autoray.py @@ -802,6 +802,117 @@ def f(x): assert y == 2 +@pytest.mark.parametrize( + "backend", + gen_params( + backends=("numpy", "torch"), + requires=("random.uniform", "zeros"), + ), +) +def test_compose_supplies_namespace(backend): + seen = [] + name = f"_test_composed_namespace_{backend}" + + @ar.compose(name=name) + def _test_composed_namespace(x, backend, namespace): + seen.append(namespace) + assert namespace._backend == backend + return namespace.zeros(1) + + x = gen_rand((1,), backend, dtype="float32") + + y = _test_composed_namespace(x) + assert ar.get_dtype_name(y) == "float32" + + y = ar.do(name, x) + assert ar.get_dtype_name(y) == "float32" + assert seen[-1] is ar.get_namespace(x) + + xp = ar.get_namespace(like=backend, dtype="float64") + y = getattr(xp, name)(x) + assert ar.get_dtype_name(y) == "float64" + assert seen[-1] is xp + + +def test_compose_namespace_from_first_argument(): + @ar.compose + def _test_composed_namespace_like(x, namespace): + return namespace.zeros(x.size) + + like = np.ones(1, dtype="float32") + + y = _test_composed_namespace_like(like) + assert y.dtype == np.dtype("float32") + + y = ar.do("_test_composed_namespace_like", like) + assert y.dtype == np.dtype("float32") + + y = ar.DoFunc("_test_composed_namespace_like")(like) + assert y.dtype == np.dtype("float32") + + +def test_compose_namespace_is_root_namespace(): + seen = [] + + @ar.compose(name="linalg._test_composed_namespace_root") + def _test_composed_namespace_root(x, namespace): + seen.append(namespace) + return x + + x = np.ones(1) + xp = ar.get_namespace(x) + y = xp.linalg._test_composed_namespace_root(x) + + assert y is x + assert seen[-1] is xp + + +def test_compose_override_not_supplied_namespace(): + @ar.compose + def _test_composed_namespace_override(x, namespace): + return namespace.sum(x) + + x = np.ones(1) + assert _test_composed_namespace_override(x) == 1.0 + + @_test_composed_namespace_override.register("numpy") + def numpy_override(x): + return "override" + + assert _test_composed_namespace_override(x) == "override" + assert ( + ar.get_namespace(x)._test_composed_namespace_override(x) == "override" + ) + + +def test_compose_namespace_falls_back_to_backend(): + @ar.compose + def _test_composed_namespace_fallback(x, namespace): + return namespace.zeros(1) + + # the first argument does not belong to the dispatched backend + y = ar.do("_test_composed_namespace_fallback", [1.0, 2.0], like="numpy") + assert y.dtype == np.dtype("float64") + + @ar.compose + def _test_composed_namespace_noargs(namespace): + return namespace.zeros(1) + + # there is no first argument at all + y = ar.do("_test_composed_namespace_noargs", like="numpy") + assert y.dtype == np.dtype("float64") + + +def test_compose_explicit_namespace_not_overridden(): + @ar.compose + def _test_composed_namespace_explicit(x, namespace): + return namespace.zeros(1) + + xp = ar.get_namespace(like="numpy", dtype="complex128") + y = _test_composed_namespace_explicit(np.ones(1), namespace=xp) + assert y.dtype == np.dtype("complex128") + + def test_builtins_complex(): re = 1.0 im = 2.0 @@ -1238,6 +1349,27 @@ def test_register_function_direct_seen_by_new_namespace(): assert ar.get_namespace(like="faux_direct").sum(x) == "custom" +def test_register_function_seen_by_existing_namespace(): + ar.autoray.register_module_alias("faux_held", "numpy") + x = np.ones(3) + + xp = ar.get_namespace(like="faux_held") + sub = xp.linalg + assert xp.sum(x) == 3.0 + assert sub.norm(x) == pytest.approx(3**0.5) + + ar.register_function("faux_held", "sum", lambda a, **kw: "custom") + ar.register_function("faux_held", "linalg.norm", lambda a, **kw: "custom") + + # namespaces are stable objects that just drop their cached lookups + assert ar.get_namespace(like="faux_held") is xp + assert xp.sum(x) == "custom" + assert xp.linalg.norm(x) == "custom" + + # including a submodule namespace already held by a caller + assert sub.norm(x) == "custom" + + def test_register_function_alias_after_import(): ar.autoray.register_module_alias("faux_alias", "numpy") x = np.array([1.0, -2.0])