WIP, ENH: Torch CPU via Array API for GFDLRegressor by sdtemple · Pull Request #129 · lanl/GFDL

sdtemple · 2026-09-14T17:39:10Z

Making changes to GFDL for compatibility with array API standards could enable the use of other namespaces (PyTorch, JAX, CuPy) and other accelerator devices, and therefore potentially speed up calculations.

This PR modifies .fit() mostly and .predict() in the base GFDL to accommodate array_api_compat functionality. The GFDLRegressor is not directly touched, as the array API modifications happen in base GFDL and invoked via super().fit() and super().predict() Test functions in test_regression.py are written to evaluate behavior with PyTorch on CPU devices.

Related to ENH: Ridge supports array API standard, we could as well? #68
Related to WIP, ENH: exact solve with array backends #115

I wrote two private utility functions in model.py to defensively ensure that the design matrix X is of a float type suitable for the device and that other arrays are converted to the namespace, device, and dtype of X. I experienced many issues in writing tests, running the modified .fit() in notebooks, and aligning with existing sklearn API conformance tests if I didn't ensure namespace, device, and (float) dtype. While these were partially developed with debugging on mps and in light of another soon-to-be PR about GFDLClassifier, the implementation and tests appear stable on CPU as the current scope. I often found it best to make most, if not all changes, in GFDL, and the other classes use that from super().fit().

The older PR has failures on sklearn API conformance tests. This often concerns edge cases about list, None, _NotAnArray inputs to the classifier, not the regressor.
https://github.com/tylerjereddy/GFDL/tree/treddy_array_api_pinv

There are ruff check issues remaining, that do not concern the changes made in this PR. ruff check --fix and manual edits were made for appropriate linting in the PR changes here.

Checklist:
[] Is there a more elegant solution than the private utility functions?
[] Is there an sklearn decorator we could use in the test suite?
[] Would some of the test functions already be captured by an sklearn decorator?
[] Benchmark speed against NumPy
[] Do the activation and weight generation functions need to be in xp? Would that affect speed?
[] Can we lower the atol in the test_torch_matches_numpy function?

AI disclosure: it wrote the scaffold of private utility functions, followed by human reviewing and editing. Test functions written by human.

- staging ground to have array api use - testing for torch cpu only right now - testing regressor only right now - test to match numpy and torch r2 score - test that predictions are in same namespace - many issues with sklearn conformance in classifier - more details to come about onehotencoder issues

- Two private utility functions for the following: - Check and/or convert all to X's dtype, namespace, & device - Ensure that X is of a float type - Test function to ensure predictions have same context as X - Test function that fitted attributes have same context as X - Test function that ints are handled gracefulyl

sdtemple · 2026-09-14T17:41:15Z

+    estimator = GFDLRegressor
+    report = r2_score


These are written to potentially be compatible with a GFDLClassifier test function or other scoring metrics in the context or a parametrization decorator mark.

sdtemple · 2026-09-14T17:42:13Z

+        (np, np.float64, np.float32),
+        (torch, torch.float64, torch.float32),
+        (np, np.float32, np.float64),
+        (torch, torch.float32, torch.float64),


This is written for stability on floating point precision 32 and 64, for potential downstream PRs on accelerator devices.

This doesn't look quite right--should be handled by upstream testing utilities/approaches. People have already thought carefully about this stuff and I don't want us to have to rehash all those discussions. Would become a nightmare to maintain this with 4-5 more namespaces.

sdtemple · 2026-09-14T17:43:57Z

+def test_int_array_api(namespace,
+                       X_dtype,
+                       y_dtype,
+                       ):
+    """Integer arrays shall be handled gracefully"""


I often ran into issues with an sklearn test with input integer arrays. This was the motivation for the private utility function _ensure_float_X, and this test is meant to confirm that that works.

sdtemple · 2026-09-14T17:46:17Z

+def test_fit_attr_context(namespace,
+                          X_dtype,
+                          y_dtype,
+                          ):


I ran into some issues where the weights, generated as NumPy arrays would in the wrong namespace, device, or dtype as X or Y. This was one of the motivations for the private utility function _check_and_convert_array. This test is meant to test that the fitted attributes are saved as the same namespace, device, and dtype of X.

I think the main point of the task was to identify what utilities/approaches sklearn uses to check these things, and use those, rather than rehashing all of that here.

In the current sklearn source I see test_logistic_regression_array_api_compliance which uses yield_namespace_device_dtype_combinations() and _array_api_for_tests()--probably not much point in me spending time here until we're matching what they've thought carefully about upstream, etc.

sdtemple · 2026-09-14T17:48:04Z

+def test_predictor_context(namespace,
+                            fit_dtype,
+                            predict_dtype,
+                            ):


This function refers to a design choice: that the .predict() calculation should be done on the device that the user specifies, i.e., the design matrix they put in. I expect reviewer discussion on this design choice. An alternative would be to make the .predict() calculations on the device of X used in .fit().

This should include a clear discussion of what sklearn does in this scenario--the main point of the task was to identify what they do and match that, rather than rehashing discussions and utilities that have already been battle tested upstream.

sdtemple · 2026-09-14T17:50:05Z

        # MoorePenrose Pseudo-Inverse, otherwise use ridge regularized form.
        if self.reg_alpha is None:
-            self.coeff_ = np.linalg.pinv(D, rtol=self.rtol) @ Y
+            self.coeff_ = xp.linalg.pinv(D, rtol=self.rtol) @ Y


This particular invocation has issues with integer arrays. I believe xp.linalg.pinv in some or all namespaces is only implemented for svd, which has an finfo() invocation. The finfo() invocation fails when given an integer array.

It would be helpful to refer to the standard itself when discussing issues. For example, if we look at: https://data-apis.org/array-api/latest/extensions/generated/array_api.linalg.pinv.html#pinv

It indicates:

x (array) – input array having shape (..., M, N) and whose innermost two dimensions form MxN matrices. Should have a floating-point data type.

So, yes, the input must be of a floating point dtype. git grep -E -i "xp.linalg.pinv" does show some example usages that might be helpful inside the sklearn source. They do seem to occasionally diverge when NumPy is used vs. not.

sdtemple · 2026-09-14T17:50:40Z

+def _check_and_convert_array(X1, X2):
+    """Convert second array to namespace, device, dtype of first if not already"""


Scaffold of this function written by AI. Reviewed and edited by human.

sdtemple · 2026-09-14T17:51:22Z

+def _ensure_float_X(xp, X):
+    """Make design matrix floating point numbers"""


Scaffold of this function written by AI. Reviewed and edited by human.

Let's just use upstream shims for this--we don't need to have AI draft this stuff when it has been battle tested upstream in SciPy and sklearn. I know SciPy has a utility to force floating for array API, etc.

sdtemple · 2026-09-14T17:52:28Z

                "scikit-learn>=1.5.0",
                "scipy>=1.13.0",
+                "array-api-compat",
+                "torch>=2.12.0,<=2.14.0",


For GFDL the max torch<=2.14.0 may not be necessary. It is on our cluster machines that I run into issues with cuda drivers not compatible with latest torch versions.

Yeah, unlikely to be a real constraint--we shouldn't have to maintain an upper bound at all. I added a more detailed comment about where these deps might live rather than dependencies. Since we're a small project, array-api-compat might be "ok" to leave here, though might make sense to pin its version.

sdtemple · 2026-09-14T18:47:00Z

The tests (both old and new) worked on the two clusters for me as well (linux-based).

tylerjereddy

I added some initial feedback--I'll expect that a lot of back and forth should happen between Emma/you/Ankan here before I'm pinged for a "real" review, but I've given you some initial feedback. The overarching message of my feedback is to avoid reinventing the wheel and match what sklearn does, falling back to what i.e., SciPy does if they are missing something.

As mentioned in person, arguably the most important component of the assigned task, as pasted below, is to determine how this is all done upstream in sklearn (preferably, since that's our target for Emma's PR) and maybe SciPy (if sklearn is missing a utility for now), so that we don't have to spend hours rehashing battle tested ways to support and test the array API from upstream.

... the source code of sklearn for how they write these tests/run them
in CI, etc.

Typically, we shouldn't really need many new tests--it should often suffice to just be able to add some decorators to existing tests to parametrize the namespaces used for their checks, and to adjust to assertions that enforce namespace/dtype preservation. That's probably a lot easier to read/review as well--take the existing tests and add a few decorators to them vs. an entire new category of tests.

tylerjereddy · 2026-09-15T17:53:30Z

                "scikit-learn>=1.5.0",
                "scipy>=1.13.0",
+                "array-api-compat",
+                "torch>=2.12.0,<=2.14.0",


I don't think any upstream example projects that I know of add these as runtime dependencies in this way. Probably best to remove and install them in the ways that match the approaches used upstream (i.e., sklearn/scipy, etc.). array-api-compat might be "ok" to leave, since we're a small project.

We can't officially turn the functionality on anyway until SciPy/sklearn turn their array API functionality on, so some manual installing on the part of the user is likely reasonable for now, but torch is something the user would just normally have to install on their own if they want to use it. And the torch binary is massive so pulling that in by default would be wasteful except in CI jobs that genuinely need it. Perhaps torch as a dependency group entry for developers for now. SciPy does have the test-array-types dependency group (https://github.com/scipy/scipy/blob/main/pyproject.toml#L131)--maybe we could use something similar for now.

I also know that SciPy uses git submodules for some of this:

e1d4eed1389f1d93318ec855730488db48475320 subprojects/array_api_compat (1.14-36-ge1d4eed) dc9e6b59061d8a116f90dcaf83c6d1878496e5ae subprojects/array_api_extra (v0.11.1)

we could consider something like that I suppose. Not sure what is best for us yet--we're a bit smaller so may have less constraints, but still pulling in torch automatically seem undesirable. array-api-compat may be "ok" to leave as a normal dep perhaps, since we're "small."

Looks like sklearn vendors some of the deps -- see their infrastructure at: sklearn/externals/array_api_compat. That's probably because adding dependencies is prohibitively expensive for them/their users. We probably don't want to do that, but we may want to do one thing that is similar--pinning our version of that lib and only updating the compatibility layer as needed, because when the array API standard changes, things may change/start to fail, so we probably want to be explicit about changes that adopt newer versions of the array API standard.

tylerjereddy · 2026-09-15T17:59:31Z

                "scikit-learn>=1.5.0",
                "scipy>=1.13.0",
+                "array-api-compat",
+                "torch>=2.12.0,<=2.14.0",


Yeah, unlikely to be a real constraint--we shouldn't have to maintain an upper bound at all. I added a more detailed comment about where these deps might live rather than dependencies. Since we're a small project, array-api-compat might be "ok" to leave here, though might make sense to pin its version.

tylerjereddy · 2026-09-15T18:20:10Z

    def fit(self, X, Y):
+        xp = array_api_compat.get_namespace(X)
+        X = _ensure_float_X(xp, X)
+        Y = _check_and_convert_array(X, Y)


Have you actually checked how sklearn does all this stuff? That was the main task assigned really--avoid reinventing the wheel and adopt their approach, especially for simplification of the process of transplanting over to Emma's sklearn PR. Looking through their PR list with a search for "array API" one can learn quite a lot about how they are doing things.

If I look at their LogisticRegression class for example, the fit() method uses this sequence:

xp, _, device = get_namespace_and_device(X) sample_weight = move_to(sample_weight, xp=xp, device=device) xp_y, _ = get_namespace(y)

see: https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/linear_model/_logistic.py#L1492-L1494

I can't think of a reason to "roll our own" approach when a large team of folks, including people on the actual standards committee, have vetted upstream approaches.

tylerjereddy · 2026-09-15T18:29:08Z

        # MoorePenrose Pseudo-Inverse, otherwise use ridge regularized form.
        if self.reg_alpha is None:
-            self.coeff_ = np.linalg.pinv(D, rtol=self.rtol) @ Y
+            self.coeff_ = xp.linalg.pinv(D, rtol=self.rtol) @ Y


It would be helpful to refer to the standard itself when discussing issues. For example, if we look at: https://data-apis.org/array-api/latest/extensions/generated/array_api.linalg.pinv.html#pinv

It indicates:

x (array) – input array having shape (..., M, N) and whose innermost two dimensions form MxN matrices. Should have a floating-point data type.

So, yes, the input must be of a floating point dtype. git grep -E -i "xp.linalg.pinv" does show some example usages that might be helpful inside the sklearn source. They do seem to occasionally diverge when NumPy is used vs. not.

tylerjereddy · 2026-09-15T18:32:48Z


    def predict(self, X):
        check_is_fitted(self)
+        xp = array_api_compat.get_namespace(X)


Here again, check sklearn source may be helpful.

The predict_proba of their LogisticRegression uses check_same_namespace(X, self, attribute="coef_", method="predict_proba") for example.

tylerjereddy · 2026-09-15T18:53:33Z

+        (np, np.float64, np.float32),
+        (torch, torch.float64, torch.float32),
+        (np, np.float32, np.float64),
+        (torch, torch.float32, torch.float64),


This doesn't look quite right--should be handled by upstream testing utilities/approaches. People have already thought carefully about this stuff and I don't want us to have to rehash all those discussions. Would become a nightmare to maintain this with 4-5 more namespaces.

tylerjereddy · 2026-09-15T18:54:54Z

+        )
+    if predict_dtype != y_pred.dtype:
+        raise ValueError(
+            ".predict output and input are not the same dtype"


we don't typically raise ValueError in tests; a lot of this looks kind of weird--you'd typically just do a plain assertion, but actually dtype passthrough/preservation is already built-in to upstream utilities so we really shouldn't spend much time reviewing or discussing things that have already been battle tested upstream for array API testing/support

tylerjereddy · 2026-09-15T19:02:26Z

+        random_state=42,
+    )
+    X = namespace.asarray(X, dtype=X_dtype, device=device)
+    y = namespace.asarray(y, dtype=y_dtype, device=device)


note that CuPy asarray() doesn't suport device=; this stuff has all been thought through upstream though.. see my other comments

tylerjereddy · 2026-09-15T19:06:08Z

+                          y_dtype,
+                          ):
+    """Fitted attributes shall be same as design matrix"""
+    os.environ["SCIPY_ARRAY_API"] = "1"


This is obviously not correct--it would bleed into other tests in our testsuite, which we don't want.

We'll need a way to turn this on/off globally for the testsuite, along with a way to select which array backends we want to test at any given time. The user/developer/CI testing could just set this and other flags externally depending on what they are trying to test.

This stuff has been hashed out upstream, so let's not reinvent the wheel.

tylerjereddy · 2026-09-15T19:11:18Z

+def test_fit_attr_context(namespace,
+                          X_dtype,
+                          y_dtype,
+                          ):


I think the main point of the task was to identify what utilities/approaches sklearn uses to check these things, and use those, rather than rehashing all of that here.

In the current sklearn source I see test_logistic_regression_array_api_compliance which uses yield_namespace_device_dtype_combinations() and _array_api_for_tests()--probably not much point in me spending time here until we're matching what they've thought carefully about upstream, etc.

Seth Temple added 2 commits September 14, 2026 11:05

sdtemple commented Sep 14, 2026

View reviewed changes

sdtemple changed the title ~~WIP: Torch CPU via Array API for GFDLRegressor~~ WIP, ENH: Torch CPU via Array API for GFDLRegressor Sep 14, 2026

sdtemple mentioned this pull request Sep 14, 2026

WIP, ENH: Torch CPU via Array API for GFDLClassifier #130

Open

sdtemple self-assigned this Sep 14, 2026

sdtemple requested a review from eiviani-lanl September 14, 2026 19:03

tylerjereddy added enhancement New feature or request array API standard support labels Sep 15, 2026

tylerjereddy requested changes Sep 15, 2026

View reviewed changes

		def _check_and_convert_array(X1, X2):
		"""Convert second array to namespace, device, dtype of first if not already"""

		def _ensure_float_X(xp, X):
		"""Make design matrix floating point numbers"""

Uh oh!

Conversation

sdtemple commented Sep 14, 2026

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!

Choose a reason for hiding this comment

Uh oh!

sdtemple commented Sep 14, 2026

Uh oh!

tylerjereddy 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants