Consider stateful API/classes, i.e. scikit-learn estimators

[seberg]

Right now, we don't have any machinery for stateful objects. The extreme case of this is scikit-learn estimators.
This isn't very important for users like SciPy who don't have much stateful API, but to push this towards scikit-learn, we would need a very clear story around this.

sklearn estimators have:

• An __init__ which sets hyperparameters only.
• Typically a fit() that gets the arguments and creates internal state.
• Other methods that use or modify the internal state

Normalizing the state is plausible, but seems hard. Typically, once we picked an implementation, we have to stick to it for all other method calls, because otherwise we cannot guarantee compatibility.¹

To make matters even more confusing, one could do something like:

with backend_opts(...):
    est = Estimator(...)

est.fit()  # do we remember the options from when it was created or here?

(Not that I think the above is very worrying either way.)

This probably means that we need a concept of "stateful methods" for such classes. If one was called, we need to remember the backend and always re-use that as well pass in the state.²
However, I need to have a clearer picture to really make up my mind about what the story is here!

There is also the component that estimators may actually make .fit() just call a single hidden _fit_internal() function to normalize things a bit first here.³

---

The next piece of the puzzle, is that scikit-learn would want a way to convert from an estimator fitted in a backend to a "normal" one (and the opposite direction).
I expect this can be handled with two stateful-methods above.

Footnotes

1. This is actually slightly different for estimators that are Array API enabled. While their state differs in which types it uses, it always has the same structure. ↩

2. One could actually see the state as part of the supported types in theory. Then plugins could effectively indicate that they "understand" another implementations state object.
   (We throw all types into one bin, but that bin could contain an estimator state object next to the arrays.) ↩

3. I would personally like to avoid this for simplicity on the scikit-learn side, but it might just be a lot clearer at least for certain estimators. ↩

[seberg]

@ev-br maybe an issue is actually easier. SciPy may have a similar issue related to JIT'ing, although I am not quite sure yet, there is Evgeni's example:

class RBFInterpolator:
     def __init__(self, y, d, neighbors: int | None, ...):
            self._xp = array_namespace(y, d)    # store the namespace
            self.neighbors = neighbors     # see __call__ below
            # ... some eager computations (similar to `.fit`)
            self._data = .... # computed from `y` and `d`

     def __call__(x):
           if self.neighbors is None:
                 return _compute(x, self._data, self._xp)  # call module level helper
           else:
                 return <something else, np-only, not jittable>

# at the module level
def _compute(x, data, xp):
    #  this is a pure function, can jit it
    # .....
    return result

Now, dispatching multiple times there could work, if the state is well defined. I am not sure it is, it isn't in the general case, so it might be best anyway to stick with expectations around statefulness with backends.

What I mean is:

class RBFinterpolator:
    @dispatchable
    def __init__(self, y, d, neighbors: int | None, ...):
        self._xp = array_namespace(y, d)    # store the namespace
        self.neighbors = neighbors     # see __call__ below

        self._data = .... # computed from `y` and `d`

    def __call__(self, x):
        return self._compute(x, self._data)

    @dispatchable
    @staticmethod
    def _compute(x, _data):
        # actual computation

Now, that would work, so long that _data is always the identical array (but it may be a different type). So we know if a backend understands that array type, it'll understand what "data" is.

---

I am not sure that the "well defined state" pattern with dispatching on all methods is great, TBH. I think it's OK, but we should figure out custom state anyway and then that may just be always the better pattern for classes.

[seberg]

@eriknw and @Schefflera-Arboricola just in case you have cycles and interest to think about this type of thing. To me the way to do statefulness is likely the most interesting big picture thing to figure out for spatch -- besides just seeing where we can use it now.
(That doesn't mean all of scikit-learn needs, mostly how we keep backend state, how we dispatch to the same backend, and how to define the dispatchable -- class decorator, decorate all methods/properties, etc.?)

[eriknw]

Thanks for continuing to push on this and writing down your thoughts @seberg! You're right, this is interesting.

Here's an idea that continues towards the explicit dispatching w.r.t. state: specify attributes of an object in @dispatchable. For example, your last example could do

    @dispatchable(["self._data", "x"])
    def __call__(self, x):
        # actual computation

which avoids a layer of indirection.

Heh, this implies a unary function can dispatch on multiple objects, such as:

    @dispatchable(["self.left", "self.right"])
    def step(self):
        ...

This may be a little complex, since the function still receives self , so it may access other data on self. This may be positive or negative depending on situation. This also doesn't prevent the pattern you shared with passing explicit stateful objects.

[seberg]

Ah, that is an interesting trick when we have well defined state on the class, we could just say which of the objects on the class are part of that "state".
We would indeed still access all of these through the original self, I don't mind that, just means it is slightly fuzzier defined.
(For the SciPy JIT use-case, it probably you probably need the indirection to a staticmethod either way, but the backend can do that also.)

There are two cases here, and I think scikit-learn needs the other one, while SciPy may actually be happy with the first:

1. There is state, but the state is fully defined. In that case, defining which of the state is important (on the dispatchable or maybe also on the class) is actually quite nice.
2. It could be that self.data here is on opaque backend specific object. In that case, no other backend can understand it. There are actually two approaches in theory:
   • We bloat the type dispatching, and just give it a custom type which we add to the dispatching system.
   • We ensure that if we dispatched once, we stick to that backend. So the step function will not dispatch again (it already knows which backend to choose).

I actually somewhat like dispatching every time, but overall, I lean towards it being impractical for complex state (but maybe not for SciPy!).

Thinking about it a bit yesterday, I was wondering about for the second case... And considering something in the direction of:

@stateful_class  # or subclass, etc.
class Estimator:
    @dispatchable(["x", "y"])
    def fit(self, x, y):
        pass
    @dispatched  # a method that doesn't make sense if not dispatched
    def other_func(self):
        pass

And then maybe the backend would be something like:

@implements_stateful_class
class EstimatorBackend:
    def __init__(self, ctx, original_estimator):
        # copy what we need from original_estimator?
        # or maybe `Estimator` can have a well defined interface
        # to call `EstimatorBackend._from_original(...)`
        # (This shouldn't have to be `__init__` to allow code-reuse
        # with user-facing code ideally, although if that is hard then nvm.)
    def fit(self, ctx, x, y):
        # backend fit.  But the __init__ will be called first
    def other_func(self, ctx):
        # might have ctx.original_estimator here also
        # context could be optional (as might allow code re-use)
        pass

[eriknw]

CC @jim22k who may find this topic interesting.

This probably means that we need a concept of "stateful methods" for such classes. If one was called, we need to remember the backend and always re-use that as well pass in the state.
[....]
how we dispatch to the same backend

A quick thought: what if we used a "simple", "low-tech" solution such as saving the preferred (default?) backend to an attribute, and @dispatchable can indicate the attribute name? This is similar to what we do in NetworkX where we dispatch on __networkx_backend__ attribute that specifies the backend name. For spatch, we would still probably prefer type-based dispatching first (if anything matches), but then look at an attribute for the default. Not sure. For stateful objects, the value of this attribute can change, so it's not limited to a class attribute.

This could look something like

    @dispatchable("self", default_backend_attribute="current_backend_name")
    def fit(self, X, y):
        ...

or whatever. I'm sure there's a better way to spell this, and maybe global config can make this easier. This could potentially interact with other dispatching policies, so I'd want to explore it more. If an API is explicitly supporting dispatching, I think it's fair to add a stateful attribute specific for dispatching.

I think @stateful_class might be convenient, but I'd want to try to make it sugar. I feel like there should be a way to specify the same behavior more explicitly by decorating each method. If you think differently, then by all means don't let me stop you and lmk how you see things!

Also, I would prefer/advise to not dispatch __init__ if it can be avoided. It gets especially complicated when there are subclasses. Dispatching on __new__ works great if it's okay to return objects from backends. I would also look into using __init_subclass__ to see if that could be used to avoid dispatching on __init__.

[seberg]

A quick thought: what if we used a "simple", "low-tech" solution such as saving the preferred (default?) backend to an attribute, and @dispatchable can indicate the attribute name?

Yeah, I think that makes a lot of sense. The question is where the backend should store its state and maybe if it isn't convenient for the backend to follow the same class structure simply?
That is why I continued a bit and wondered if it makes sense to store a backend class instance (which is automatically dispatched already).

Also, I would prefer/advise to not dispatch init if it can be avoided. It gets especially complicated when there are subclasses.

Hmmm, I hadn't thought about subclasses! (I think you mentioned it at some point). If a dispatching class is subclassed things might get awkward. I do wonder if that is really relevant to things like scikit-learn estimators (everyone subclasses their base-classes, but I don't really see much use in subclassing the specific estimator -- but I guess someone might; it might not be too bad, if all they do is use public API to add other things, but...).

[ev-br]

Thanks @seberg . Not that I understand details here, I'll just note that at least for SciPy it would be best if we could somehow avoid the dispatching work (and the corresponding perf penalty) on each __call__ invocation. Currently, we can store the namespace on the instance and reuse it at each __call__.

[seberg]

@ev-br I was thinking about this a bit and I think the SciPy use-case here should maybe really just be solved by dispatching __new__? And that already should "just work".
The backend should just implement its own version as a subclass. It'll be slightly weird that you get a subclass back from __new__, but that would be the main awkwardness. (If we really dislike, we have to dispatch every time anyway although that is a bit faster if you already know where to dispatch to.)

The thing about the sklearn one is that it needs a few more things:

1. Dispatching can only happen delayed (after __new__/__init__).
2. There needs to be state, and that state may not be rather backend dependent.
3. There seems to be a wish to have a way to convert at least back to "normal sklearn state".

[seberg]

I have a small try in gh-41, with the idea of dispatching once on certain methods. And then general methods and properties just forward.
That is based on a class decorator right now that replaces methods/properties with such forwarding versions.

(Conversion can be added as a to_sklearn() or so method always I think. Or also properties or a function to get the desired state on a per estimator level of course.)