scikit-learn / scikit-learn/scikit-learn

Expose `get_namespace_and_device` and `move_estimator_to` in public API

Open
#34,135 8 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

API Array API
Dominant language
Python
Stars
67.3k
Forks
27.4k
Avg merge
1d 15h
Merged PRs (30d)
58

Description

[!WARNING]
This issue is not yet ready for a PR. If you are interested in contributing to scikit-learn, please have a look at our contributing guidelines, and in particular the sections for new contributors and the "Needs triage" label.

I just stumbled over the error message of the new ValueError added in #33076:

ValueError: Inputs passed to LinearRegression.predict() must use the same namespace and the same device as those passed to fit(). Array namespaces used during fit (sklearn.externals.array_api_compat.numpy) and predict (sklearn.externals.array_api_compat.torch) differ. You can move the estimator to the same namespace and device as X with: 'from sklearn.utils._array_api import move_estimator_to; xp, _, device = get_namespace_and_device(X); estimator = move_estimator_to(estimator, xp, device)'

For instance triggered by:

import sklearn
import numpy as np
import torch
from sklearn.linear_model import LinearRegression

rng = np.random.default_rng(0)
X = rng.normal(size=(10, 5))
y = rng.normal(size=10)

reg = LinearRegression().fit(X, y)
X_xp = torch.asarray(X)
reg.predict(X_xp)

with sklearn.config_context(array_api_dispatch=True):
    reg.predict(X_xp)

The recommended action is to run:

from sklearn.utils._array_api import move_estimator_to
xp, _, device = get_namespace_and_device(X)
estimator = move_estimator_to(estimator, xp, device)

Both methods are also mentioned in https://scikit-learn.org/dev/modules/array_api.html.

I am wondering why move_estimator_to and get_namespace_and_device are not exposed in the public API then. Did we forget to do that or was that a conscious choice? (I tried to find discussion on that but couldn't find any.)

If we expose both methods in the API, users can access the docstrings from the website and import via from sklearn.utils import move_estimator_to, get_namespace_and_device.
Also, we may want to expose get_namespace and get_device, though I don't find that so important.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with sklearn.utils._array_api and the public sklearn.utils import surface to understand how get_namespace_and_device and move_estimator_to are currently defined and exposed. Read the array API documentation, especially its references to these methods, then confirm the intended scope with maintainers; done means a decided public API change with matching exports and documentation, or a documented reason not to expose them.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
api, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.