WIP, ENH, TST: Array API support initial draft by eiviani-lanl · Pull Request #2 · eiviani-lanl/scikit-learn

eiviani-lanl · 2026-10-15T01:59:21Z

This is the first draft of Array API support. In making these changes, I had to take some artistic liberties with regard to partial_fit, as there aren't many good examples of sklearn estimators that have both partial_fit and Array API support.

I think we should also consider how closely we expect the NumPy and Torch paths to agree numerically, particularly for float32 inputs. I made some changes to the solvers in our models (Ridge(solver="svd") and storing A and B in float64) to improve consistency between the two paths, but I'm not sure if this is what users will expect. sklearn seems to accept the numerical differences produced by Ridge, so maybe these changes aren't strictly necessary.

sdtemple · 2026-10-15T19:41:36Z

Thanks for putting this initial draft up. We noticed you opened this during our group meeting and some of us (especially me) will start reviewing this.

sdtemple

I did an initial pass of comments only while looking at the code. There are some notes on things I'll circle back to and verify if that code is necessary or could be edited.

sdtemple · 2026-10-21T16:27:18Z

+    xp, _ = get_namespace(X)
+    tmp = -X + xp.min(X, axis=1)[:, xp.newaxis]
+    X[...] = xp.exp(tmp)
+    X /= xp.sum(X, axis=1)[:, xp.newaxis]


Can this just be return inplace_softmax(-X)? I guess perhaps it cannot because it is an in place call?

sdtemple · 2026-10-21T16:32:42Z

+    xp, _ = get_namespace(X)
+    X[...] = xp.maximum(X, xp.zeros_like(X))


I know at lanl/GFDL#130 that I used get_namespace_and_device so that we could put zeros on the device.

xp, _, device = get_namespace_and_device(z) zero = xp.asarray(0, dtype=z.dtype, device=device) return xp.maximum(zero, z)

Need to check that there is not a lack of device issue

sdtemple · 2026-10-21T16:35:49Z

+@parametrize_with_checks(
+    estimators=[ExtremeLearningClassifier(), ExtremeLearningRegressor()],
+    expected_failed_checks=get_expected_failed_checks,
+)
+def test_sklearn_array_api_conformance(estimator, check):
+    # None of the sklearn models require this test, maybe we can pull it out?
+    with config_context(array_api_dispatch=True):
+        check(estimator)


I believe check(estimator) may already handle the with config_context(array_api_dispatch=True) for any array API tests when SCIPY_ARRAY_API=1 is turned on. This would negate the need to write the expected failure of _NotAnArray because it would only do that test for base NumPy.

I'll check the source code.

I actually think we can remove this entire test. I believe that check_array_api_input is captured in _yield_array_api_checks which is captured in _yield_checks. And within the array API related tests in sklearn.utils.estimator_checks they are repeatedly using with config_context(...): and turning those on when an estimator has the array API tags.

Okay, great.

sdtemple · 2026-10-21T16:38:05Z

+        ExtremeLearningClassifier(ridge_alpha=0.1),
+        ExtremeLearningRegressor(ridge_alpha=0.1),
+    ],
+)


At sdtemple/GFDL (scikit-learn#130 of lanl/GFDL), I test a lot more combinations of weights, activations, and layer sizes.

I'll check the sklearn source code to see if that is necessary or if it is only checking input validation-related things.

sdtemple · 2026-10-21T17:00:12Z

+        xp, _, device_ = get_namespace_and_device(
+            X,
+            y,
+        )


I started to drop the y argument to this function. I think the convention is that everything should be put on X namespace and device. I have to verify if adding y when y is a NumPy array and X is a torch tensor would result in xp as NumPy.

I have actually been hoping that people keep y on a NumPy array for downstream OneHotEncoder, though I think we should cast y to NumPy before OneHotEncoder anyway as a defensive mechanism.

If you run the following, you'll see that NumPy is the xp brought in:

import torch import numpy from sklearn.utils._array_api import get_namespace x = numpy.zeros(5) y = torch.zeros(5) get_namespace(x,y)

If you check the implementation in LogisticRegression, they only get the namespace from X.

sdtemple · 2026-10-21T17:01:42Z

+                self.coef_ = self.coef_.ravel()
+            else:
+                self.coef_ = _ravel(self.coef_)
        return self


Is the ravel necessary?

sdtemple · 2026-10-21T17:03:44Z

        # shape: (n_samples, n_features)
-        X, y = validate_data(self, X, y)
+        xp, _, device_ = get_namespace_and_device(X)
+        y_xp, _, _ = get_namespace_and_device(y)


I have to check if this is strictly necessary.

sdtemple · 2026-10-21T17:03:50Z

+            self,
+            X,
+            y,
+            dtype=[xp.float64, xp.float32],


I have to check if this is strictly necessary.

This I believe achieves what I do with:

if not xp.isdtype(X.dtype, "real floating"): default_float_dtype = _max_precision_float_dtype(xp, device) X = xp.astype(X, default_float_dtype)

When I delete the above from sdtemple/GFDL and put in your version of validate_data it works as desired. They don't always do that dtype=[xp.float64, xp.float32] with validate_data, but they do it in .fit() for logistic regression and the GMM, so let's keep it.

sdtemple · 2026-10-21T17:04:07Z

+            # Keep classes_ in y's original namespace and device.
+            self.classes_ = y_xp.sort(y_xp.unique_values(y))
+
+            matches = y_xp.expand_dims(y, axis=1) == y_xp.expand_dims(
+                self.classes_, axis=0
+            )


I have to check if this is strictly necessary.

sdtemple · 2026-10-21T20:46:27Z

        # design matrix shape: (n_samples, sum_hidden+n_features)
        # or (n_samples, sum_hidden)
-        D = self._design_matrix(X)
+        D = self._design_matrix(X, xp)


I think it'd be better practice to have these methods like _design_matrix read the xp from X instead of passing xp.

sdtemple · 2026-10-21T21:24:38Z

+        return self._predict_proba(X, xp)

-    def _predict_proba(self, X):
+    def _predict_proba(self, X, xp):


The convention in @_logistic.py seems to be that you pass in only X and use get_namespace within the function.

sdtemple

Here are more comments about the test functions. I'll be making edits on my branch with respect to those comments.

sdtemple · 2026-10-21T21:50:52Z

+    # many fail otherwise
+    xp, device = _array_api_for_tests(array_namespace, device_name, dtype_name)
+    X_np, y_np = generator(
+        n_samples=10000,


Do you think we need these many samples? The sklearn utility, that we thought had too few, used 30. The logistic regression implementation uses iris data (which has 150? samples).

sdtemple · 2026-10-21T21:51:36Z

+        assert get_namespace(pred_xp)[0] == get_namespace(X_xp)[0]
+        assert array_device(pred_xp) == array_device(X_xp)
+        assert pred_xp.dtype == X_xp.dtype
+


Do you think these aren't already covered by the sklearn array API compliance checks?

sdtemple · 2026-10-21T21:53:20Z

+        assert_allclose(
+            move_to(attr_xp, xp=np, device="cpu"), attr_np, rtol=rtol, atol=atol
+        )
+        assert attr_xp.dtype == X_xp.dtype


We can systematically test this in a for loop. This is what sklearn does. Although that is a for loop only over "coef_", it would be good to follow their code patterns and its sustainable if other attributes get added.

While I do evidence that test_logistic_regression_array_api_compliance checks the coefficients, they do this for 1 small dataset (iris). That's the brittle part of the test.

sdtemple · 2026-10-21T22:07:00Z

+        check_array_api_input_and_values(
+            name,
+            estimator,
+            array_namespace,
+            device_name=device_name,
+            dtype_name=dtype_name,
+        )


I actually think we can remove this entire test. I believe that check_array_api_input is captured in _yield_array_api_checks which is captured in _yield_checks. And within the array API related tests in sklearn.utils.estimator_checks they are repeatedly using with config_context(...): and turning those on when an estimator has the array API tags.

sdtemple · 2026-10-21T22:07:14Z

+@parametrize_with_checks(
+    estimators=[ExtremeLearningClassifier(), ExtremeLearningRegressor()],
+    expected_failed_checks=get_expected_failed_checks,
+)
+def test_sklearn_array_api_conformance(estimator, check):
+    # None of the sklearn models require this test, maybe we can pull it out?
+    with config_context(array_api_dispatch=True):
+        check(estimator)


I actually think we can remove this entire test. I believe that check_array_api_input is captured in _yield_array_api_checks which is captured in _yield_checks. And within the array API related tests in sklearn.utils.estimator_checks they are repeatedly using with config_context(...): and turning those on when an estimator has the array API tags.

sdtemple · 2026-10-21T22:07:47Z

+        )
+
+
 @parametrize_with_checks([ExtremeLearningClassifier(), ExtremeLearningRegressor()])


Would it make sense to add a version of the estimators with the pinv turned on as well?

I think we can get by with just this test, which will cover 2 of your other tests, as long as we put a specific @skip_if_array_api_compat_not_configured decorator on the test function.

Why do we call this conformance instead of compliance?

sdtemple · 2026-10-21T22:09:46Z

+        assert get_namespace(pred_xp)[0] == get_namespace(X_xp)[0]
+        assert array_device(pred_xp) == array_device(X_xp)
+        assert pred_xp.dtype == X_xp.dtype
+


I think these are captured in the array API compliance checks already.

sdtemple · 2026-10-21T22:11:17Z

+def test_array_api_activations(
+    activation,
+    weight_init,
+    estimator,
+    generator,
+    array_namespace,
+    device_name,
+    dtype_name,
+):


So, is this just a test that the activation functions behave with a smaller dataset? I feel like there is a lot of redundancy. And, don't the activations get used elsewhere in the neural_network module. We may need to consider that.

I would say in this case we may want to synthesize some of the tests into 1.

sdtemple · 2026-10-22T23:05:53Z

+
+@pytest.mark.parametrize("ridge_alpha", [None, 2])
+@pytest.mark.parametrize("hidden_layer_sizes", [(10,), (10, 10)])
+@pytest.mark.parametrize("batch", [100, 1000])


I don't see much value in exploring 2 batch sizes

sdtemple · 2026-10-22T23:16:59Z

+            self.A = xp.zeros(
+                (D.shape[1], D.shape[1]), dtype=xp.float64, device=device_
+            )
+            self.B = xp.zeros(
+                (D.shape[1], y.shape[1]), dtype=xp.float64, device=device_
+            )


These xp.zeros are problematic because you set dtype=xp.float64. This crashes for me when I run with PYTORCH_ENABLE_MPS_FALLBACK=1

sdtemple · 2026-10-22T23:22:01Z

+        assert get_namespace(pred_xp)[0] == get_namespace(X_xp)[0]
+        assert array_device(pred_xp) == array_device(X_xp)
+        assert pred_xp.dtype == X_xp.dtype
+
+        assert_allclose(
+            move_to(pred_xp, xp=np, device="cpu"),
+            pred_np,
+            rtol=rtol,
+            atol=atol,
+        )
+
+        attr_xp = getattr(elm_xp, "coef_")
+        attr_np = getattr(elm_np, "coef_")
+
+        assert_allclose(
+            move_to(attr_xp, xp=np, device="cpu"), attr_np, rtol=rtol, atol=atol
+        )
+        assert attr_xp.dtype == X_xp.dtype
+        assert array_device(attr_xp) == array_device(X_xp)


These lines don't have to be in the with block for the array API dispatch. In general, only .fit(), .predict, and .partial_fit need to be. Also, some of these assert statements I think are already captured by the compliance check.

I would suggest we remove the coefficient matching for partial_fit. There isn't much precedent for that.

sdtemple · 2026-10-22T23:42:43Z

+        self.coef_ = xp.astype(self.coef_, X.dtype)
        if self.coef_.ndim == 2 and self.coef_.shape[1] == 1:
-            self.coef_ = self.coef_.ravel()
+            self.coef_ = _ravel(self.coef_)


Under this if guard, I think xp.squeeze(self.coef_, axis=1) achieves the same without using the private utility. I'll make that change in my branch, as I did in lanl/GFDL

sdtemple · 2026-10-22T23:43:43Z

        if self.coef_.ndim == 2 and self.coef_.shape[1] == 1:
-            self.coef_ = self.coef_.ravel()
+            if _is_numpy_namespace(xp):
+                self.coef_ = self.coef_.ravel()
+            else:
+                self.coef_ = _ravel(self.coef_)


Under this if guard, I think xp.squeeze(self.coef_, axis=1) achieves the same without using the private utility. I'll make this change in my branch, as I did in the lanl/GFDL PR

sdtemple · 2026-10-23T16:47:29Z

+        if not hasattr(self, "classes_"):
+            if _is_numpy_namespace(y_xp):
+                classes = move_to(classes, xp=y_xp, device=y_device_)
+                classes = y_xp.asarray(classes, device=y_device_)
+                self._label_binarizer = LabelBinarizer().fit(classes)
+                self.classes_ = self._label_binarizer.classes_
+
+            else:
+                if y.ndim == 2 and y.shape[1] == 1:
+                    y = y_xp.reshape(y, (-1,))
+                classes = move_to(classes, xp=y_xp, device=y_device_)
+                self.classes_ = y_xp.sort(y_xp.unique_values(classes))
+


@eiviani-lanl Is there some reason related to backward compatibility why we have all these if and else blocks based on _is_numpy_namspace?

LabelBinarizer has Array API support. https://scikit-learn.org/stable/modules/array_api.html#estimators

It is possible we could write this more succinctly.

I don't see instances of the Naive Bayes estimator of LogisticRegression having these if and else blocks for handling NumPy label binarization differently. I would say we can use LabelBinarizer for all array namespaces, as this has the Array API support.

I proposed such a change on my branch. It is fewer code lines and clearer without if statements.

sdtemple · 2026-10-23T16:54:54Z

+            W_xp, _, W_device_ = get_namespace_and_device(self.W_[0])
+
+            if (xp != W_xp) or (device_ != W_device_):
+                raise ValueError(
+                    "All batches must use same device and namespace , %r, "
+                    "but got %r." % ((W_xp, W_device_), (xp, device_))
+                )
+            dtype = self.W_[0].dtype


This is pretty verbose. However, the suggested edits from GPT Enterprise lead to default errors about namespace and device mismatches, which don't seem as informative as your custom ValueError here. I lean towards keeping it.

sdtemple · 2026-10-23T17:52:09Z

+        y = column_or_1d(y, warn=True)
+        check_classification_targets(y)
+
+        y_xp, _, y_device_ = get_namespace_and_device(y)


The convention from another Array API estimator like Naive Bayes is to use xp_y. That implies that this is the namespace of y, not y on some namespace xp. I made such naming changes for y, classes, and W

sdtemple · 2026-10-23T19:00:26Z

@eiviani-lanl I have gone through a few rounds now and I am satisfied with the initial review comments I gave. I made proposed code changes to address those comments at my branch elm_array_api_sdtemple. I do not want to merge my branch to this PR without your consent and your feedback on the edits. You can run

git diff elm_array_api..elm_array_api_sdtemple

to see how I primarily simplified you if blocks with _is_numpy_namespace, shortened redundant tests, and followed some naming convention aligned with logistic regression and Naive Bayes Array API estimators.

You can also use

git diff main..elm_array_api_sdtemple

to see the overall changes from the sklearn PR?

The main squabble I still have is with the partial fit function for the Array API support. It has a lot of duplicated code and can be brittle for atol and rtol values in float32. I do feel like partial fit is already covered pretty extensively in other tests for the NumPy namespace and the analogous Array API estimator GaussianNB with partial_fit does not check backend numerical precision of partial_fit.

* SI disclosure: I did not use SI, except I did use GPT 6.1 Sol when reviewing Emma's work * There were some Array API styles that Emma used in her sklearn PR that I preferred over mine. I also identified some redundant tests in my sklearn PR review. Here I am porting over those suggested changes. * eiviani-lanl/scikit-learn#2 * Use `xp.asarray` instead of `move_to` and `xp.astype` * Use `dtype=[xp.float64, xp.float32]` in `validate_data` * Only `xp.squeeze` when needed. Fixes a bug with classifier which will go upstream when I rebase lanl#130 * Removed redundant tests which were identified during review

sdtemple · 2026-10-26T18:37:57Z

+        assert array_device(pred_xp) == array_device(X_xp)
+        assert pred_xp.dtype == X_xp.dtype
+
+        assert_allclose(


array-api-extra has a utility xpx.testing.assert_close that handles the dtype precision and also checks namespace and device. That package is vendored in scikit-learn. I think we should use it so that we don't manually set rtol.

I'll post an update to my branch.

sdtemple · 2026-10-26T18:40:17Z

+):
+    xp, device = _array_api_for_tests(array_namespace, device_name, dtype_name)
+    X_np, y_np = generator(
+        n_samples=10000,


I would suggest we do a very small dataset for partial_fit, e.g., 6 samples and a batch size of 2. There isn't much precedent for such a large dataset in partial_fit checks.

* SI disclosure: I did not use SI, except I did use GPT 6.1 Sol when reviewing Emma's work * There were some Array API styles that Emma used in her sklearn PR that I preferred over mine. I also identified some redundant tests in my sklearn PR review. Here I am porting over those suggested changes. * eiviani-lanl/scikit-learn#2 * Use `xp.asarray` instead of `move_to` and `xp.astype` * Use `dtype=[xp.float64, xp.float32]` in `validate_data` * Only `xp.squeeze` when needed. Fixes a bug with classifier which will go upstream when I rebase lanl#130 * Removed redundant tests which were identified during review

WIP, ENH, TST: Array API support initial draft

97aced4

eiviani-lanl changed the base branch from main to extreme-learning-machine October 15, 2026 02:07

sdtemple reviewed Oct 21, 2026

View reviewed changes

sdtemple reviewed Oct 22, 2026

View reviewed changes

sdtemple reviewed Oct 23, 2026

View reviewed changes

sdtemple mentioned this pull request Oct 23, 2026

ENH: Array API for GFDLRegressor lanl/GFDL#129

Open

sdtemple reviewed Oct 26, 2026

View reviewed changes

		xp, _ = get_namespace(X)
		X[...] = xp.maximum(X, xp.zeros_like(X))

		)


		@parametrize_with_checks([ExtremeLearningClassifier(), ExtremeLearningRegressor()])

Conversation

eiviani-lanl commented Oct 15, 2026

Uh oh!

sdtemple commented Oct 15, 2026

Uh oh!

sdtemple left a comment

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!

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!

sdtemple left a comment

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!

sdtemple Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

sdtemple Oct 22, 2026 •

edited

Loading