some new features

This commit is contained in:
ilgazca
2025-07-30 17:09:11 +03:00
parent db5d46760a
commit 8019bd3b7c
20616 changed files with 4375466 additions and 8 deletions

View File

@ -0,0 +1,29 @@
BSD 3-Clause License
Copyright (c) 2018, Quansight-Labs
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
* Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
* Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

View File

@ -0,0 +1,116 @@
"""
.. note:
If you are looking for overrides for NumPy-specific methods, see the
documentation for :obj:`unumpy`. This page explains how to write
back-ends and multimethods.
``uarray`` is built around a back-end protocol, and overridable multimethods.
It is necessary to define multimethods for back-ends to be able to override them.
See the documentation of :obj:`generate_multimethod` on how to write multimethods.
Let's start with the simplest:
``__ua_domain__`` defines the back-end *domain*. The domain consists of period-
separated string consisting of the modules you extend plus the submodule. For
example, if a submodule ``module2.submodule`` extends ``module1``
(i.e., it exposes dispatchables marked as types available in ``module1``),
then the domain string should be ``"module1.module2.submodule"``.
For the purpose of this demonstration, we'll be creating an object and setting
its attributes directly. However, note that you can use a module or your own type
as a backend as well.
>>> class Backend: pass
>>> be = Backend()
>>> be.__ua_domain__ = "ua_examples"
It might be useful at this point to sidetrack to the documentation of
:obj:`generate_multimethod` to find out how to generate a multimethod
overridable by :obj:`uarray`. Needless to say, writing a backend and
creating multimethods are mostly orthogonal activities, and knowing
one doesn't necessarily require knowledge of the other, although it
is certainly helpful. We expect core API designers/specifiers to write the
multimethods, and implementors to override them. But, as is often the case,
similar people write both.
Without further ado, here's an example multimethod:
>>> import uarray as ua
>>> from uarray import Dispatchable
>>> def override_me(a, b):
... return Dispatchable(a, int),
>>> def override_replacer(args, kwargs, dispatchables):
... return (dispatchables[0], args[1]), {}
>>> overridden_me = ua.generate_multimethod(
... override_me, override_replacer, "ua_examples"
... )
Next comes the part about overriding the multimethod. This requires
the ``__ua_function__`` protocol, and the ``__ua_convert__``
protocol. The ``__ua_function__`` protocol has the signature
``(method, args, kwargs)`` where ``method`` is the passed
multimethod, ``args``/``kwargs`` specify the arguments and ``dispatchables``
is the list of converted dispatchables passed in.
>>> def __ua_function__(method, args, kwargs):
... return method.__name__, args, kwargs
>>> be.__ua_function__ = __ua_function__
The other protocol of interest is the ``__ua_convert__`` protocol. It has the
signature ``(dispatchables, coerce)``. When ``coerce`` is ``False``, conversion
between the formats should ideally be an ``O(1)`` operation, but it means that
no memory copying should be involved, only views of the existing data.
>>> def __ua_convert__(dispatchables, coerce):
... for d in dispatchables:
... if d.type is int:
... if coerce and d.coercible:
... yield str(d.value)
... else:
... yield d.value
>>> be.__ua_convert__ = __ua_convert__
Now that we have defined the backend, the next thing to do is to call the multimethod.
>>> with ua.set_backend(be):
... overridden_me(1, "2")
('override_me', (1, '2'), {})
Note that the marked type has no effect on the actual type of the passed object.
We can also coerce the type of the input.
>>> with ua.set_backend(be, coerce=True):
... overridden_me(1, "2")
... overridden_me(1.0, "2")
('override_me', ('1', '2'), {})
('override_me', ('1.0', '2'), {})
Another feature is that if you remove ``__ua_convert__``, the arguments are not
converted at all and it's up to the backend to handle that.
>>> del be.__ua_convert__
>>> with ua.set_backend(be):
... overridden_me(1, "2")
('override_me', (1, '2'), {})
You also have the option to return ``NotImplemented``, in which case processing moves on
to the next back-end, which in this case, doesn't exist. The same applies to
``__ua_convert__``.
>>> be.__ua_function__ = lambda *a, **kw: NotImplemented
>>> with ua.set_backend(be):
... overridden_me(1, "2")
Traceback (most recent call last):
...
uarray.BackendNotImplementedError: ...
The last possibility is if we don't have ``__ua_convert__``, in which case the job is
left up to ``__ua_function__``, but putting things back into arrays after conversion
will not be possible.
"""
from ._backend import *
__version__ = '0.8.8.dev0+aa94c5a4.scipy'

View File

@ -0,0 +1,704 @@
import typing
import types
import inspect
import functools
from . import _uarray
import copyreg
import pickle
import contextlib
from ._uarray import ( # type: ignore
BackendNotImplementedError,
_Function,
_SkipBackendContext,
_SetBackendContext,
_BackendState,
)
__all__ = [
"set_backend",
"set_global_backend",
"skip_backend",
"register_backend",
"determine_backend",
"determine_backend_multi",
"clear_backends",
"create_multimethod",
"generate_multimethod",
"_Function",
"BackendNotImplementedError",
"Dispatchable",
"wrap_single_convertor",
"wrap_single_convertor_instance",
"all_of_type",
"mark_as",
"set_state",
"get_state",
"reset_state",
"_BackendState",
"_SkipBackendContext",
"_SetBackendContext",
]
ArgumentExtractorType = typing.Callable[..., tuple["Dispatchable", ...]]
ArgumentReplacerType = typing.Callable[
[tuple, dict, tuple], tuple[tuple, dict]
]
def unpickle_function(mod_name, qname, self_):
import importlib
try:
module = importlib.import_module(mod_name)
qname = qname.split(".")
func = module
for q in qname:
func = getattr(func, q)
if self_ is not None:
func = types.MethodType(func, self_)
return func
except (ImportError, AttributeError) as e:
from pickle import UnpicklingError
raise UnpicklingError from e
def pickle_function(func):
mod_name = getattr(func, "__module__", None)
qname = getattr(func, "__qualname__", None)
self_ = getattr(func, "__self__", None)
try:
test = unpickle_function(mod_name, qname, self_)
except pickle.UnpicklingError:
test = None
if test is not func:
raise pickle.PicklingError(
f"Can't pickle {func}: it's not the same object as {test}"
)
return unpickle_function, (mod_name, qname, self_)
def pickle_state(state):
return _uarray._BackendState._unpickle, state._pickle()
def pickle_set_backend_context(ctx):
return _SetBackendContext, ctx._pickle()
def pickle_skip_backend_context(ctx):
return _SkipBackendContext, ctx._pickle()
copyreg.pickle(_Function, pickle_function)
copyreg.pickle(_uarray._BackendState, pickle_state)
copyreg.pickle(_SetBackendContext, pickle_set_backend_context)
copyreg.pickle(_SkipBackendContext, pickle_skip_backend_context)
def get_state():
"""
Returns an opaque object containing the current state of all the backends.
Can be used for synchronization between threads/processes.
See Also
--------
set_state
Sets the state returned by this function.
"""
return _uarray.get_state()
@contextlib.contextmanager
def reset_state():
"""
Returns a context manager that resets all state once exited.
See Also
--------
set_state
Context manager that sets the backend state.
get_state
Gets a state to be set by this context manager.
"""
with set_state(get_state()):
yield
@contextlib.contextmanager
def set_state(state):
"""
A context manager that sets the state of the backends to one returned by :obj:`get_state`.
See Also
--------
get_state
Gets a state to be set by this context manager.
""" # noqa: E501
old_state = get_state()
_uarray.set_state(state)
try:
yield
finally:
_uarray.set_state(old_state, True)
def create_multimethod(*args, **kwargs):
"""
Creates a decorator for generating multimethods.
This function creates a decorator that can be used with an argument
extractor in order to generate a multimethod. Other than for the
argument extractor, all arguments are passed on to
:obj:`generate_multimethod`.
See Also
--------
generate_multimethod
Generates a multimethod.
"""
def wrapper(a):
return generate_multimethod(a, *args, **kwargs)
return wrapper
def generate_multimethod(
argument_extractor: ArgumentExtractorType,
argument_replacer: ArgumentReplacerType,
domain: str,
default: typing.Optional[typing.Callable] = None,
):
"""
Generates a multimethod.
Parameters
----------
argument_extractor : ArgumentExtractorType
A callable which extracts the dispatchable arguments. Extracted arguments
should be marked by the :obj:`Dispatchable` class. It has the same signature
as the desired multimethod.
argument_replacer : ArgumentReplacerType
A callable with the signature (args, kwargs, dispatchables), which should also
return an (args, kwargs) pair with the dispatchables replaced inside the
args/kwargs.
domain : str
A string value indicating the domain of this multimethod.
default: Optional[Callable], optional
The default implementation of this multimethod, where ``None`` (the default)
specifies there is no default implementation.
Examples
--------
In this example, ``a`` is to be dispatched over, so we return it, while marking it
as an ``int``.
The trailing comma is needed because the args have to be returned as an iterable.
>>> def override_me(a, b):
... return Dispatchable(a, int),
Next, we define the argument replacer that replaces the dispatchables inside
args/kwargs with the supplied ones.
>>> def override_replacer(args, kwargs, dispatchables):
... return (dispatchables[0], args[1]), {}
Next, we define the multimethod.
>>> overridden_me = generate_multimethod(
... override_me, override_replacer, "ua_examples"
... )
Notice that there's no default implementation, unless you supply one.
>>> overridden_me(1, "a")
Traceback (most recent call last):
...
uarray.BackendNotImplementedError: ...
>>> overridden_me2 = generate_multimethod(
... override_me, override_replacer, "ua_examples", default=lambda x, y: (x, y)
... )
>>> overridden_me2(1, "a")
(1, 'a')
See Also
--------
uarray
See the module documentation for how to override the method by creating
backends.
"""
kw_defaults, arg_defaults, opts = get_defaults(argument_extractor)
ua_func = _Function(
argument_extractor,
argument_replacer,
domain,
arg_defaults,
kw_defaults,
default,
)
return functools.update_wrapper(ua_func, argument_extractor)
def set_backend(backend, coerce=False, only=False):
"""
A context manager that sets the preferred backend.
Parameters
----------
backend
The backend to set.
coerce
Whether or not to coerce to a specific backend's types. Implies ``only``.
only
Whether or not this should be the last backend to try.
See Also
--------
skip_backend: A context manager that allows skipping of backends.
set_global_backend: Set a single, global backend for a domain.
"""
try:
return backend.__ua_cache__["set", coerce, only]
except AttributeError:
backend.__ua_cache__ = {}
except KeyError:
pass
ctx = _SetBackendContext(backend, coerce, only)
backend.__ua_cache__["set", coerce, only] = ctx
return ctx
def skip_backend(backend):
"""
A context manager that allows one to skip a given backend from processing
entirely. This allows one to use another backend's code in a library that
is also a consumer of the same backend.
Parameters
----------
backend
The backend to skip.
See Also
--------
set_backend: A context manager that allows setting of backends.
set_global_backend: Set a single, global backend for a domain.
"""
try:
return backend.__ua_cache__["skip"]
except AttributeError:
backend.__ua_cache__ = {}
except KeyError:
pass
ctx = _SkipBackendContext(backend)
backend.__ua_cache__["skip"] = ctx
return ctx
def get_defaults(f):
sig = inspect.signature(f)
kw_defaults = {}
arg_defaults = []
opts = set()
for k, v in sig.parameters.items():
if v.default is not inspect.Parameter.empty:
kw_defaults[k] = v.default
if v.kind in (
inspect.Parameter.POSITIONAL_ONLY,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
):
arg_defaults.append(v.default)
opts.add(k)
return kw_defaults, tuple(arg_defaults), opts
def set_global_backend(backend, coerce=False, only=False, *, try_last=False):
"""
This utility method replaces the default backend for permanent use. It
will be tried in the list of backends automatically, unless the
``only`` flag is set on a backend. This will be the first tried
backend outside the :obj:`set_backend` context manager.
Note that this method is not thread-safe.
.. warning::
We caution library authors against using this function in
their code. We do *not* support this use-case. This function
is meant to be used only by users themselves, or by a reference
implementation, if one exists.
Parameters
----------
backend
The backend to register.
coerce : bool
Whether to coerce input types when trying this backend.
only : bool
If ``True``, no more backends will be tried if this fails.
Implied by ``coerce=True``.
try_last : bool
If ``True``, the global backend is tried after registered backends.
See Also
--------
set_backend: A context manager that allows setting of backends.
skip_backend: A context manager that allows skipping of backends.
"""
_uarray.set_global_backend(backend, coerce, only, try_last)
def register_backend(backend):
"""
This utility method sets registers backend for permanent use. It
will be tried in the list of backends automatically, unless the
``only`` flag is set on a backend.
Note that this method is not thread-safe.
Parameters
----------
backend
The backend to register.
"""
_uarray.register_backend(backend)
def clear_backends(domain, registered=True, globals=False):
"""
This utility method clears registered backends.
.. warning::
We caution library authors against using this function in
their code. We do *not* support this use-case. This function
is meant to be used only by users themselves.
.. warning::
Do NOT use this method inside a multimethod call, or the
program is likely to crash.
Parameters
----------
domain : Optional[str]
The domain for which to de-register backends. ``None`` means
de-register for all domains.
registered : bool
Whether or not to clear registered backends. See :obj:`register_backend`.
globals : bool
Whether or not to clear global backends. See :obj:`set_global_backend`.
See Also
--------
register_backend : Register a backend globally.
set_global_backend : Set a global backend.
"""
_uarray.clear_backends(domain, registered, globals)
class Dispatchable:
"""
A utility class which marks an argument with a specific dispatch type.
Attributes
----------
value
The value of the Dispatchable.
type
The type of the Dispatchable.
Examples
--------
>>> x = Dispatchable(1, str)
>>> x
<Dispatchable: type=<class 'str'>, value=1>
See Also
--------
all_of_type
Marks all unmarked parameters of a function.
mark_as
Allows one to create a utility function to mark as a given type.
"""
def __init__(self, value, dispatch_type, coercible=True):
self.value = value
self.type = dispatch_type
self.coercible = coercible
def __getitem__(self, index):
return (self.type, self.value)[index]
def __str__(self):
return f"<{type(self).__name__}: type={self.type!r}, value={self.value!r}>"
__repr__ = __str__
def mark_as(dispatch_type):
"""
Creates a utility function to mark something as a specific type.
Examples
--------
>>> mark_int = mark_as(int)
>>> mark_int(1)
<Dispatchable: type=<class 'int'>, value=1>
"""
return functools.partial(Dispatchable, dispatch_type=dispatch_type)
def all_of_type(arg_type):
"""
Marks all unmarked arguments as a given type.
Examples
--------
>>> @all_of_type(str)
... def f(a, b):
... return a, Dispatchable(b, int)
>>> f('a', 1)
(<Dispatchable: type=<class 'str'>, value='a'>,
<Dispatchable: type=<class 'int'>, value=1>)
"""
def outer(func):
@functools.wraps(func)
def inner(*args, **kwargs):
extracted_args = func(*args, **kwargs)
return tuple(
Dispatchable(arg, arg_type)
if not isinstance(arg, Dispatchable)
else arg
for arg in extracted_args
)
return inner
return outer
def wrap_single_convertor(convert_single):
"""
Wraps a ``__ua_convert__`` defined for a single element to all elements.
If any of them return ``NotImplemented``, the operation is assumed to be
undefined.
Accepts a signature of (value, type, coerce).
"""
@functools.wraps(convert_single)
def __ua_convert__(dispatchables, coerce):
converted = []
for d in dispatchables:
c = convert_single(d.value, d.type, coerce and d.coercible)
if c is NotImplemented:
return NotImplemented
converted.append(c)
return converted
return __ua_convert__
def wrap_single_convertor_instance(convert_single):
"""
Wraps a ``__ua_convert__`` defined for a single element to all elements.
If any of them return ``NotImplemented``, the operation is assumed to be
undefined.
Accepts a signature of (value, type, coerce).
"""
@functools.wraps(convert_single)
def __ua_convert__(self, dispatchables, coerce):
converted = []
for d in dispatchables:
c = convert_single(self, d.value, d.type, coerce and d.coercible)
if c is NotImplemented:
return NotImplemented
converted.append(c)
return converted
return __ua_convert__
def determine_backend(value, dispatch_type, *, domain, only=True, coerce=False):
"""Set the backend to the first active backend that supports ``value``
This is useful for functions that call multimethods without any dispatchable
arguments. You can use :func:`determine_backend` to ensure the same backend
is used everywhere in a block of multimethod calls.
Parameters
----------
value
The value being tested
dispatch_type
The dispatch type associated with ``value``, aka
":ref:`marking <MarkingGlossary>`".
domain: string
The domain to query for backends and set.
coerce: bool
Whether or not to allow coercion to the backend's types. Implies ``only``.
only: bool
Whether or not this should be the last backend to try.
See Also
--------
set_backend: For when you know which backend to set
Notes
-----
Support is determined by the ``__ua_convert__`` protocol. Backends not
supporting the type must return ``NotImplemented`` from their
``__ua_convert__`` if they don't support input of that type.
Examples
--------
Suppose we have two backends ``BackendA`` and ``BackendB`` each supporting
different types, ``TypeA`` and ``TypeB``. Neither supporting the other type:
>>> with ua.set_backend(ex.BackendA):
... ex.call_multimethod(ex.TypeB(), ex.TypeB())
Traceback (most recent call last):
...
uarray.BackendNotImplementedError: ...
Now consider a multimethod that creates a new object of ``TypeA``, or
``TypeB`` depending on the active backend.
>>> with ua.set_backend(ex.BackendA), ua.set_backend(ex.BackendB):
... res = ex.creation_multimethod()
... ex.call_multimethod(res, ex.TypeA())
Traceback (most recent call last):
...
uarray.BackendNotImplementedError: ...
``res`` is an object of ``TypeB`` because ``BackendB`` is set in the
innermost with statement. So, ``call_multimethod`` fails since the types
don't match.
Instead, we need to first find a backend suitable for all of our objects.
>>> with ua.set_backend(ex.BackendA), ua.set_backend(ex.BackendB):
... x = ex.TypeA()
... with ua.determine_backend(x, "mark", domain="ua_examples"):
... res = ex.creation_multimethod()
... ex.call_multimethod(res, x)
TypeA
"""
dispatchables = (Dispatchable(value, dispatch_type, coerce),)
backend = _uarray.determine_backend(domain, dispatchables, coerce)
return set_backend(backend, coerce=coerce, only=only)
def determine_backend_multi(
dispatchables, *, domain, only=True, coerce=False, **kwargs
):
"""Set a backend supporting all ``dispatchables``
This is useful for functions that call multimethods without any dispatchable
arguments. You can use :func:`determine_backend_multi` to ensure the same
backend is used everywhere in a block of multimethod calls involving
multiple arrays.
Parameters
----------
dispatchables: Sequence[Union[uarray.Dispatchable, Any]]
The dispatchables that must be supported
domain: string
The domain to query for backends and set.
coerce: bool
Whether or not to allow coercion to the backend's types. Implies ``only``.
only: bool
Whether or not this should be the last backend to try.
dispatch_type: Optional[Any]
The default dispatch type associated with ``dispatchables``, aka
":ref:`marking <MarkingGlossary>`".
See Also
--------
determine_backend: For a single dispatch value
set_backend: For when you know which backend to set
Notes
-----
Support is determined by the ``__ua_convert__`` protocol. Backends not
supporting the type must return ``NotImplemented`` from their
``__ua_convert__`` if they don't support input of that type.
Examples
--------
:func:`determine_backend` allows the backend to be set from a single
object. :func:`determine_backend_multi` allows multiple objects to be
checked simultaneously for support in the backend. Suppose we have a
``BackendAB`` which supports ``TypeA`` and ``TypeB`` in the same call,
and a ``BackendBC`` that doesn't support ``TypeA``.
>>> with ua.set_backend(ex.BackendAB), ua.set_backend(ex.BackendBC):
... a, b = ex.TypeA(), ex.TypeB()
... with ua.determine_backend_multi(
... [ua.Dispatchable(a, "mark"), ua.Dispatchable(b, "mark")],
... domain="ua_examples"
... ):
... res = ex.creation_multimethod()
... ex.call_multimethod(res, a, b)
TypeA
This won't call ``BackendBC`` because it doesn't support ``TypeA``.
We can also use leave out the ``ua.Dispatchable`` if we specify the
default ``dispatch_type`` for the ``dispatchables`` argument.
>>> with ua.set_backend(ex.BackendAB), ua.set_backend(ex.BackendBC):
... a, b = ex.TypeA(), ex.TypeB()
... with ua.determine_backend_multi(
... [a, b], dispatch_type="mark", domain="ua_examples"
... ):
... res = ex.creation_multimethod()
... ex.call_multimethod(res, a, b)
TypeA
"""
if "dispatch_type" in kwargs:
disp_type = kwargs.pop("dispatch_type")
dispatchables = tuple(
d if isinstance(d, Dispatchable) else Dispatchable(d, disp_type)
for d in dispatchables
)
else:
dispatchables = tuple(dispatchables)
if not all(isinstance(d, Dispatchable) for d in dispatchables):
raise TypeError("dispatchables must be instances of uarray.Dispatchable")
if len(kwargs) != 0:
raise TypeError(f"Received unexpected keyword arguments: {kwargs}")
backend = _uarray.determine_backend(domain, dispatchables, coerce)
return set_backend(backend, coerce=coerce, only=only)