DOC Clarify how mixed array input types handled in array api #31452

lucyleeow · 2025-05-29T10:54:47Z

Reference Issues/PRs

Update documentation now that we have reach consensus on these issues:

closes #31286 (I think no more work is required for this as currently all converted metrics are outputting array in the same namespace/device as input)

Towards #28668
Towards #31274

I have not said close for these 2 because we would need to implement conversion of mixed inputs.

What does this implement/fix? Explain your changes.

Clarify in docs:

everything follows X for estimators
everything follows y_pred for scoring functions
when a scoring function outputs several values, an array of the same namespace and device is output

I realise that this update to the docs is not yet true, but I think this okay as I don't think this will get into 1.7, thus will be in dev docs only, and I think we will have implemented it by next release.
Otherwise, happy to leave this PR and implement the mixed input type conversion first.

Any other comments?

I tried to give reasons as to why we made these decisions. Hopefully it is clear and concise (and correct...), but happy to change/fix.

cc @ogrisel @betatim maybe @StefanieSenger

github-actions · 2025-05-29T10:55:44Z

✔️ Linting Passed

All linting checks passed. Your pull request is in excellent shape! ☀️

_{Generated for commit: 5833e9a. Link to the linter CI: here}

StefanieSenger

Thanks for your work @lucyleeow!

I like the clarity of how this new addition is written and I am sure it will help many users. I have left a few minor comments, and I have some more general comments.

What about starting this section with a fast to grasp takeaway for our user's mental model, that we have taken months to uncover: To clearly state that other array inputs follow X and y_pred in terms of array library and device?
Another feedback would be to re-think the heading of this paragraph. It now not only talks about return types, but also on the internal processing of inputs that follow X and y_pred. Users interested in that might not expect this information under this heading.
And, there is something about a sub-part that I do not understand (and which might also happen to users): In lines 194-212 you showcase an example, where Xs namespace is changed during the process of going through a pipeline and there are several things unclear to me:

When a pipeline starts out on user input with X and y both being on CPU, then I would not expect scikit-learn to take any measures on automatically processing these on GPU for intermediate steps, since we don't even know if users have GPU...
At first sight it seems that the y follows X principle here is broken and upon thinking about it, I could then see that this part is to demonstrate how y being on the different namespace on the beginning of a new step of a pipeline demonstrates exactly how it follows X without that the user needs to take care of it. But right now, there is only a very subtle hint on this in this paragraph and this idea is not clearly expressed.
When you write "Note that scikit-learn pipelines do not allow transformation of y (to avoid
:ref:leakage <data_leakage>).": does that even include changes in which namespace y is stored, is that the reason you bring this up here?
I wonder, why X and y both need to be on CPU for TargetEncoder? Is there a technical reason or a more goal-oriented reason in terms of performance?

Oh my, that has gotten to become a lot of feedback. 😬 (sorry)

doc/modules/array_api.rst

StefanieSenger · 2025-05-29T11:45:20Z

doc/modules/array_api.rst

+The `predict` and `transform` method subsequently expect
 inputs from the same array library and device as the data passed to the `fit`
 method.


Maybe mention that as of now, scikit-learn does not interfere with how array libraries handle arrays from a different namespace or device passed to predict or transform. Something like:

"If arrays from a different library or on a different device are passed, behavior depends on the array library: it may raise an error or silently convert them."

Sorry I am confused here, why would an array library be doing things within our predict and transform methods? 🤔

When arrays from different namespaces are passed into a function, then the namespace that the function comes from determines the handling. For instance in xp.add(array1, array2 ) the namespace of xp could be torch and array1 is a torch array, but array2 is not.Then torch will handle this or raise an error.

Okay you've raised some good points 🤔 I'm going to raise these back in the issue

doc/modules/array_api.rst

lucyleeow · 2025-05-30T06:20:07Z

doc/modules/array_api.rst

+:class:`~sklearn.linear_model.Ridge`. `y` cannot be transformed by the pipeline
+(recall scikit-learn pipelines do not allow transformation of `y`) but as
+:class:`~sklearn.linear_model.Ridge` is able to accept mixed input types,
+this is not a problem and the pipeline is able to be run.


At first sight it seems that the y follows X principle here is broken and upon thinking about it, I could then see that this part is to demonstrate how y being on the different namespace on the beginning of a new step of a pipeline demonstrates exactly how it follows X without that the user needs to take care of it. But right now, there is only a very subtle hint on this in this paragraph and this idea is not clearly expressed.

@StefanieSenger this is the only part that explicit talks about this. I tried to make it clearer but open to suggests to improve. Thanks!

lucyleeow · 2025-05-30T06:21:29Z

These are very good suggestions. I have tried to address them all.

then I would not expect scikit-learn to take any measures on automatically processing these on GPU for intermediate steps, since we don't even know if users have GPU...

Sorry I have made it more explicit now. The function transformer step is: FunctionTransformer(partial(torch.asarray, device="cuda:0")), - which explicitly moves X to GPU.
For reference I was mostly copying Oliviers example here: #31274 (comment)

I wonder, why X and y both need to be on CPU for TargetEncoder? Is there a technical reason or a more goal-oriented reason in terms of performance?

In my mind, the categorical (string) data in X was being target encoded to numbers in TargetEncoder so X has to be on CPU (as due to it containing string data), and y "needs" to be on the same device as X for the estimator to work (TargetEncoder also uses y) - though if we are doing everything follows X, y technically doesn't need to be on CPU. I will try to make things clearer. I was worried I was going into too much detail and making things too long 🤔

When you write "Note that scikit-learn pipelines do not allow transformation of y (to avoid
:ref:leakage <data_leakage>).": does that even include changes in which namespace y is stored, is that the reason you bring this up here?

Yes, pipelines can't touch y.

betatim · 2025-06-02T13:07:17Z

doc/modules/array_api.rst

+Estimators and scoring functions are able to accept input arrays
+from different array libraries and/or devices. When mixed input arrays are
+passed, scikit-learn silently converted to make them all consistent.


I've just realised that this means that I can/should/could call some_estimator.fit(X_torch_on_gpu0, y_numpy, sample_weights=weights_cupy_gpu1) or some other more obscure combination. Feels like mayhem, but I guess that is "fine"? Maybe somewhere in a guide we can mention that this is technically allowed but that maybe people shouldn't do this?

Suggestion (copied and slightly adjusted from the numpy docs):

.. warning:: While the mixing of array types from different array libraries may be convenient, it is not recommended. It might have unexpected behavior in corner cases. Users should prefer explicitly passing matching types whenever possible.

betatim · 2025-06-02T13:09:02Z

doc/modules/array_api.rst

+Estimators and scoring functions are able to accept input arrays
+from different array libraries and/or devices. When mixed input arrays are
+passed, scikit-learn silently converted to make them all consistent.


Suggested change

Estimators and scoring functions are able to accept input arrays

from different array libraries and/or devices. When mixed input arrays are

passed, scikit-learn silently converted to make them all consistent.

Estimators and scoring functions are able to accept input arrays

from different array libraries and/or devices. When a mixed set of input arrays is

passed, scikit-learn converts arrays as needed to make them all consistent.

What do you think?

betatim · 2025-06-02T13:09:34Z

doc/modules/array_api.rst

+passed, scikit-learn silently converted to make them all consistent.
+
+For estimators, the rule is **"everything follows `X`"** - mixed array inputs are
+converted so they are all match the array library and device of `X`.


Suggested change

converted so they are all match the array library and device of `X`.

converted so that they all match the array library and device of `X`.

betatim · 2025-06-02T13:10:01Z

doc/modules/array_api.rst

+For estimators, the rule is **"everything follows `X`"** - mixed array inputs are
+converted so they are all match the array library and device of `X`.
+For scoring functions the rule is **"everything follows `y_pred`"** - mixed array
+inputs are converted so they are all match the array library and device of `y_pred`.


Suggested change

inputs are converted so they are all match the array library and device of `y_pred`.

inputs are converted so that they all match the array library and device of `y_pred`.

betatim · 2025-06-02T13:11:41Z

doc/modules/array_api.rst

+For scoring functions the rule is **"everything follows `y_pred`"** - mixed array
+inputs are converted so they are all match the array library and device of `y_pred`.
+
+When a function or method has been called with Array API compatible inputs, the
 convention is to return array values of the same array container type and


How about using "array container type" as a way to make it clear that we are talking about numpy, cupy, pytorch, etc arrays and not "an array of float32s"?

The sentences above are already quite long, so I'm not sure if we should do this.

I find the expression "array container type" (which is already in the sentence) quite vague. In other parts of the docs we have talked about "array library" and "array namespace", so maybe like that (?):

Suggested change

convention is to return array values of the same array container type and

convention is to return a type from the same array library and

betatim · 2025-06-02T13:13:21Z

doc/modules/array_api.rst

+
+When an estimator is fitted with an Array API compatible `X`, all other
+array inputs, including constructor arguments, (e.g., `y`, `sample_weight`,
+`weights_init`) will be silently converted to match the array library and


Suggested change

`weights_init`) will be silently converted to match the array library and

`weights_init`) will be converted to match the array library and

"silently" makes me worried we will have the "explicit is better than implicit" police coming round to complain. And I can convince myself that words ending in -ly are not needed anyway :D

betatim · 2025-06-02T13:19:52Z

doc/modules/array_api.rst

+This allows estimators to accept mixed input types, enabling `X` to be moved
+to a different device within a pipeline, without explicitly moving `y`.
+Note that scikit-learn pipelines do not allow transformation of `y` (to avoid
+:ref:`leakage <data_leakage>`).


Suggested change

This allows estimators to accept mixed input types, enabling `X` to be moved

to a different device within a pipeline, without explicitly moving `y`.

Note that scikit-learn pipelines do not allow transformation of `y` (to avoid

:ref:`leakage <data_leakage>`).

This behaviour enables pipelines to switch from processing on

the CPU to processing on the GPU at a specific point in the pipeline.

This is still a bit clunky :-/ I think it is enough to mention the use-case and then show the example below, which contains a longer explanation/reminder about the fact that you can't move y.

betatim · 2025-06-02T13:21:19Z

doc/modules/array_api.rst

+  `partial(torch.asarray, device="cuda")`, which moves `X` to GPU, to improve
+  performance in the next step.
+* :class:`~sklearn.linear_model.Ridge`, whose performance can be improved when
+  passed arrays on GPU, as it only takes numerical inputs.


"on GPU", "on a GPU", "on the GPU"? Not sure which one is right/sounds best (same applies to "on CPU" above) - do you have a preference?

I have a slight preference for "on a GPU".

betatim · 2025-06-02T13:25:42Z

doc/modules/array_api.rst

+~~~~~~~~~~~~~~~~~
+
+When an Array API compatible `y_pred` is passed to a scoring function,
+all other array inputs (e.g., `y_true`, `sample_weight`) will be silently converted


Suggested change

all other array inputs (e.g., `y_true`, `sample_weight`) will be silently converted

all other array inputs (e.g., `y_true`, `sample_weight`) will be converted

ogrisel

Some partial feedback.

ogrisel · 2025-06-02T15:32:57Z

doc/modules/array_api.rst

+For scoring functions the rule is **"everything follows `y_pred`"** - mixed array
+inputs are converted so they are all match the array library and device of `y_pred`.
+
+When a function or method has been called with Array API compatible inputs, the


Nitpick: I cannot find the discussion anymore, but I think the official casing style for the spec name is "array API" (or the "Python array API") instead of "Array API".

Suggested change

When a function or method has been called with Array API compatible inputs, the

When a function or method has been called with array API compatible inputs, the

At least it's the casing style of the main header of the official spec site: https://data-apis.org/array-api/latest/

ogrisel · 2025-06-02T15:38:09Z

doc/modules/array_api.rst

+  `partial(torch.asarray, device="cuda")`, which moves `X` to GPU, to improve
+  performance in the next step.
+* :class:`~sklearn.linear_model.Ridge`, whose performance can be improved when
+  passed arrays on GPU, as it only takes numerical inputs.


I have a slight preference for "on a GPU".

StefanieSenger

These are cool improvements, thank you @lucyleeow!
Upon reading through, I came across some minor things, that could be improved.

StefanieSenger · 2025-06-03T07:11:32Z

doc/modules/array_api.rst

+For scoring functions the rule is **"everything follows `y_pred`"** - mixed array
+inputs are converted so they are all match the array library and device of `y_pred`.
+
+When a function or method has been called with Array API compatible inputs, the
 convention is to return array values of the same array container type and


I find the expression "array container type" (which is already in the sentence) quite vague. In other parts of the docs we have talked about "array library" and "array namespace", so maybe like that (?):

Suggested change

convention is to return array values of the same array container type and

convention is to return a type from the same array library and

StefanieSenger · 2025-06-03T07:13:03Z

doc/modules/array_api.rst

-Similarly, when an estimator is fitted with Array API compatible inputs, the
-fitted attributes will be arrays from the same library as the input and stored
-on the same device. The `predict` and `transform` method subsequently expect
+For details and reasons behind this choice, see below.


I believe that is not really needed here.

Suggested change

For details and reasons behind this choice, see below.

StefanieSenger · 2025-06-03T07:25:59Z

doc/modules/array_api.rst

+* :class:`~sklearn.preprocessing.FunctionTransformer`, with the function
+  `partial(torch.asarray, device="cuda")`, which moves `X` to GPU, to improve


I am not sure if it renders correctly, but in my mind that's a bit more obvious (and spares people to look up the usage of FunctionTransformer before they can read on:

Suggested change

* :class:`~sklearn.preprocessing.FunctionTransformer`, with the function

`partial(torch.asarray, device="cuda")`, which moves `X` to GPU, to improve

* :class:`~sklearn.preprocessing.FunctionTransformer`

(func=partial(torch.asarray, device="cuda")), which moves `X` to GPU, to improve

StefanieSenger · 2025-06-03T07:29:47Z

doc/modules/array_api.rst

+  `partial(torch.asarray, device="cuda")`, which moves `X` to GPU, to improve
+  performance in the next step.
+* :class:`~sklearn.linear_model.Ridge`, whose performance can be improved when
+  passed arrays on GPU, as it only takes numerical inputs.


On "as it only takes numerical inputs": that is not making for improved performance in GPU compared to CPU. We could instead name an actual reason (some matrix computation going on within Ridge) or leave this out completely?

StefanieSenger · 2025-06-03T07:37:32Z

doc/modules/array_api.rst

+When an Array API compatible `y_pred` is passed to a scoring function,
+all other array inputs (e.g., `y_true`, `sample_weight`) will be silently converted
+to match the array library and device of `y_pred`, if they do not already.
+This allows scoring functions to accept mixed input types, which enables it to be


Suggested change

This allows scoring functions to accept mixed input types, which enables it to be

This allows scoring functions to accept mixed input types, which enables them to be

lucyleeow added 2 commits May 29, 2025 15:54

doc mixed input

24d8f1a

wording

e2fc189

lucyleeow added the Array API label May 29, 2025

github-actions bot added the Documentation label May 29, 2025

lucyleeow mentioned this pull request May 29, 2025

Add array API support to median_absolute_error #31406

Merged

StefanieSenger reviewed May 29, 2025

View reviewed changes

review

5833e9a

lucyleeow commented May 30, 2025

View reviewed changes

lucyleeow mentioned this pull request Jun 2, 2025

Automatically move y (and sample_weight) to the same device and namespace as X #28668

Open

betatim reviewed Jun 2, 2025

View reviewed changes

ogrisel reviewed Jun 2, 2025

View reviewed changes

StefanieSenger reviewed Jun 3, 2025

View reviewed changes

	converted so they are all match the array library and device of `X`.
	converted so that they all match the array library and device of `X`.

	inputs are converted so they are all match the array library and device of `y_pred`.
	inputs are converted so that they all match the array library and device of `y_pred`.

	convention is to return array values of the same array container type and
	convention is to return a type from the same array library and

	`weights_init`) will be silently converted to match the array library and
	`weights_init`) will be converted to match the array library and

	all other array inputs (e.g., `y_true`, `sample_weight`) will be silently converted
	all other array inputs (e.g., `y_true`, `sample_weight`) will be converted

	When a function or method has been called with Array API compatible inputs, the
	When a function or method has been called with array API compatible inputs, the

		* :class:`~sklearn.preprocessing.FunctionTransformer`, with the function
		`partial(torch.asarray, device="cuda")`, which moves `X` to GPU, to improve

	This allows scoring functions to accept mixed input types, which enables it to be
	This allows scoring functions to accept mixed input types, which enables them to be

Uh oh!

DOC Clarify how mixed array input types handled in array api #31452

Are you sure you want to change the base?

DOC Clarify how mixed array input types handled in array api #31452

Uh oh!

Conversation

Uh oh!

Reference Issues/PRs

What does this implement/fix? Explain your changes.

Any other comments?

Uh oh!

Uh oh!

✔️ Linting Passed

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!