Figure out SciPy use-case

[seberg]

@lucascolley mentioned that SciPy has currently a light-weight dispatching that is limited to three specific types though (numpy, cupy, torch). Implemented across:

• https://github.com/scipy/scipy/blob/main/scipy/signal/_signal_api.py
• https://github.com/scipy/scipy/blob/main/scipy/signal/_support_alternative_backends.py
• https://github.com/scipy/scipy/blob/main/scipy/signal/_delegators.py

I think that everything for this use-case (practically) in-place, with maybe slight additions and some points on how to achieve things.
(Plus, one thing I realized is that the way we document backends may be very much open and likely needs customization for libraries.)

The basic pattern here would probably be (untested):

import abc

class JaxLike(abc.ABC):
    # We could just match "~jax.numpy:ndarray" below as well.
    # this is (maybe ridiculously so) closer to the original
    # by going via `__array_namespace__`.  (N.B. we look it up
    # on the class)
    @classmethod
    def __subclasscheck__(cls, other):
        get_xp = getattr(other, "__array_namespace__")
        if get_xp is None:
            return False

        # Ideally, we could just do this, but JAX doesn't define
        # this as a class method, but in spatch we prefer to match
        # types.  (Otherwise, we would just have to match anything
        # with `__array_function__` and use `should_run`...
        try:
           xp = get_xp()
        except Exception:
           pass
        else:
           return is_jax(xp)

        # Just guess that if it came from jax a jax module
        # we can use the JAX version.
        return getattr(xp, "__module__", "").startswith("jax.")


class JaxBackend:
    name = "jax_array_api"
    require_opt_in = False  # pure type dispatcher
    # Note that the "@" indicates that JaxLike is an abstract class
    primary_types = [f"@{JaxLike.__module__}:JaxLike"]
    secondary_types = []  # none, we don't allow NumPy coercion.
    functions = {
        # trivial to generate, you can do so even on import.
        "scipy.submodule:name": {
            "function": "jax.submodule:name",
            # more options, we don't need right now.
        }
    }


from spatch.backend_system import BackendSystem

_bs = BackendSystem(
    None,  # don't load entry-points for now (no 3rd party backends)
    "_SCIPY_INTERNAL_BACKENDS",  # spatch env-var prefix
    # could also allow subclasses for ndarray with "~numpy:ndarray".
    # Doesn't really matter, since we'll fall back to it anyway:
    default_primary_types=["numpy:ndarray"],
    # Register SciPy internal backends explicitly:
    backends=[JaxBackend],
)


@_bs.dispatchable(["argname"])
def scipy_function(argname, other_arg):
    # SciPy (numpy/fallback) version.

The JaxLike is a bit much, because spatch insists (or strongly prefers) to use type based semantics, but the Array API/JAX don't allow us to match "jax array API like" just based on the type (which could be changed there for sure), but this'll work in practice (or just change to list JAX types explicitly...).

Now, you are using also NumPy's __array_function__ pattern to have an ..._signature function that extracts the "relevant arguments", while the above passes a list of argument names instead.
The dedicated function version will be a lot faster for now (unless we move things to C, it may be a 1/3 or higher difference!).
I think we need to allow a callable there anyway if just for complex cases (i.e. numpy.concatenate considers all entries in the first argument as "relevant args", it needs to unpack that sequence.)

So @lucascolley if you are curious to test things out and need e.g. the version where you can do @_bs.dispatchable(..._signature), let's just add it here quickly (it is very easy of course).

[lucascolley]

I am curious to test things out, if you are at that stage!

Note that there is also https://github.com/scipy/scipy/blob/main/scipy/ndimage/_support_alternative_backends.py and https://github.com/scipy/scipy/blob/main/scipy/special/_support_alternative_backends.py

[seberg]

The above should work (of course this is alpha). I would think we can fold ndimage/special in just as well. Just need more backends?
So, I suspect you can make things work, but I am sure you will find some things you want to change or make easier.

One annoyance is filling in functions. It may be nicer to only do it when ndimage is imported (just to keep the information localized), while the backend itself I require creating immediately.

That works fine though, you can just extend the functions dict (besides possible future user inspection API). It just must be filled in when the @dispatchable is created.

One thing I just saw that (at least in ndimage?) The "fallback" mode actually still tries to fetch __array_namespace__, but it then wraps the NumPy version.
Since that is a single "wrapper" (the "dask" version may be similar for chunking), that is something that I am not trying to make it easy yet. You would have to do this "manually".

But of course, manually, just can be a trick like this:

# Module (or class) to "manually" auto-wrap/try xp to numpy version

def __getattr__(self, name):
    original = getattr(..., name)
    def wrap(*args, **kwargs):
        xp = ...
        res = original(*args, **kwargs)
        # code to re-wrap from NumPy result

    return wrap  # plus functools.wraps logic

Then, add the function as {"original_mode:original_name": {"function": "this.module:original_name"}}

And register this just with:

class ArrayAPIEnabled(abc.ABC):
    @classmethod
    def __subclasscheck__(cls, other):
        # or just `hasattr`
        return getattr(other, "__array_namespace__", None) is not None

If you do that, you would have to make sure either add ArrayAPIEnabled as the last one. (Users of entry-points, would have to do more than that, but for backends= we can partially use the passed order.)

[seberg]

Ah, I now see that there is one more thing that you need! xp = ... doesn't work yet. For that you will need to use "uses_context": True (in the functions dict). And I'll have to add that information to the context, which I'll just do right now :).

EDIT: Need to re-org things slightly, so let's see if we want this. I always suspected we will, but had considered it a lower priority, it's just... you already have the situation.

[seberg]

OK, you'll want this: #13 and add the uses_context: True for all functions :).
I suspect you are good with that (next thing may be the dispatch argument extraction via a function).

[seberg]

OK, and gh-14 contains also the parameter extraction part from a function. (I also realized that the old version assumes parameters are the first positional args unless you pass a dict, but that wasn't documented -- I fixed that but via inspect.signature, so not sure if we'll keep it).

Either way, I hope that based on that (or once it's merged), you may have everything to get something working.