Backend API#

Backends have to do the heavy lifting of using spatch. At the moment we suggest to check the example.

Entry point definition#

To extend an existing library with a backend, you need to define a Python entry-point. This entry point includes the necessary information for spatch to find and dispatch to your backend.

Before writing a backend, you need to think about a few things:

  • Which types do you accept? This could be NumPy, dask, jax, etc. arrays. There are two kind of types “primary” and “secondary”. Secondary types are types that you accept, but intend to convert (e.g. you may accept NumPy arrays, but always convert them to cupy). Primary types are types that you work with and return. If you have more than one primary type, you may need to take a lot of care about which type to return!

  • Do you change behavior of the library (or existing major backends) by e.g. providing a faster but less accurate implementation? In that case, your backend should likely only be used if prioritized by the user.

Please check the example linked above. These example entry-points include code that means running them modifies them in-place if the @implements decorator is used (see next section).

Some of the most important things are:

### name

The name of the backend, must match the name of the entry-point.

### primary_types and secondary_types

Primary and secondary types are defined by a sequence of strings and stored as primary_types and secondary_types attributes in the entry-point. As everything, they are identified by "__module__:__qualname__" strings, for example "numpy:ndarray".

We suggest that backends initially use exact type matches. At least for arrays, subclasses commonly change behavior in large ways (breaking Liskov substitution principle), so functions may not behave correctly for them anyway. However, we do support the following, e.g.:

  • "numpy:ndarray" to match NumPy arrays exactly

  • "~numpy:ndarray" to match any subclass of NumPy arrays

  • "@module:qualname" to match any subclass of an abstract base class

If you use an abstract base class, note that you must take a lot of care:

  • The abstract base class must be cheap to import, because we cannot avoid importing it.

  • Since we can’t import all classes, spatch has no ability to order abstract classes correctly (but we order them last if a primary type, which is typically right).

  • spatch will not guarantee correct behavior if an ABC is mutated at runtime.

### requires_opt_in

A boolean indicating whether your backend should be active by default. Typically, set this to True for a type dispatching backend and False otherwise. A backend doing both needs to set it to True, and must use should_run with a context to disambiguate.

Based on library guidance and feedback a non-type dispatching backend may also set this to True if it’s behavior matches the library behavior closely.

Warning

Always check with library guidelines or reach out to authors before setting this to True.

The golden rule is that behavior must be the same if your backend is installed but nothing else changes.

And remember that identical is typically impossible in numerical computing. So always check with the library (or other backends) authors first what they consider to be an acceptable difference.

Failure to do this, will just mean that spatch needs a way to block backends or only allow specific ones, and then everyone loses…

### functions

A mapping of library functions to your implementations. All fields use the __module__:__qualname__ identifiers to avoid immediate import. The following fields are supported for each function:

  • function: The implementation to dispatch to.

  • should_run (optional): A function that gets all inputs (and context) and can decide to defer. Unless you know things will error, try to make sure that this function is light-weight.

  • uses_context: Whether the implementation needs a DispatchContext.

  • additional_docs (optional): Brief text to add to the documentation of the original function. We suggest keeping this short but including a link, but the library guidance should be followed.

spatch provides tooling to help create this mapping.

### Manual backend prioritization

spatch tries to order backends based on the types, but this cannot get the order right always. Thus, you can manually prioritize your backend over others by defining for example higher_priority_than = ["default"] or lower_priority_than = ["default"].

It is your responsibility to ensure that these prioritizations make sense and are acceptable to other backends.

### More?

Note

Missing information? We are probably missing important information currently. For example, maybe a version? (I don’t think we generally need it, but it may be interesting.)

Implementations for dispatchable functions#

Spatch provides a decorator to mark functions as implementations. Applying this decorator is not enough to actually dispatch but it can be used to fill in the necessary information into the entry-point definition.

class spatch.backend_utils.BackendImplementation(backend_name: str)#
implements(api_identity: str | Callable, *, should_run: str | Callable | None = None, uses_context: bool = False)#

Mark function as an implementation of a dispatchable library function.

This is a decorator to facilitate writing an entry-point and potentially attaching should_run. It is not necessary to use this decorator and you may want to wrap it into a convenience helper for your purposes.

Parameters:

  • api_identity – The original function to be wrapped. Either as a string identity or the callable (from which the identifier is extracted).

  • should_run –

    Callable or a string with the module and name for a should_run function. A should_run function takes a DispatchInfo object as first argument and otherwise all arguments of the wrapped function. It must return True if the backend should be used.

    It is the backend author’s responsibility to ensure that should_run is not called unnecessarily and is ideally very light-weight (if it doesn’t return True). (Any return value except True is considered falsy to allow the return to be used for diagnostics in the future.)

  • uses_context – Whether the function should be passed a DispatchContext object as first argument.

set_should_run(backend_func: str | Callable)#

Alternative decorator to set the should_run function.

This is a decorator to decorate a function as should_run. If the function is a lambda or called _ it is attached to the wrapped backend function (cannot be a string). Otherwise, it is assumed that the callable can be found at runtime.

Parameters:

backend_func – The backend function for which we want to set should_run.

class spatch.backend_system.DispatchContext(types: tuple[type], name: str, _state: tuple)#

Additional information passed to backends about the dispatching.

DispatchContext is passed as first (additional) argument to should_run``and to a backend implementation (if desired). Some backends will require the ``types attribute.

types#

The (unique) types we are dispatching for. It is possible that not all types are passed as arguments if the user is requesting a specific type.

Backends that have more than one primary types must use this information to decide which type to return. I.e. if you allow mixing of types and there is more than one type here, then you have to decide which one to return (promotion of types). E.g. a cupy.ndarray and numpy.ndarray together should return a cupy.ndarray.

Backends that strictly match a single primary type can safely ignore this (they always return the same type).

Note

This is a frozenset currently, but please consider it a sequence.

Type:

Sequence[type]

prioritized#

Whether the backend is prioritized. You may use this for example when deciding if you want to defer with should_run. Or it may be fine to use this to decide that e.g. a NumPy array will be converted to a cupy array, but only if prioritized.

Type:

bool