[sdtemple]

This PR modifies .fit(), partial_fit() and .predict() in the base GFDL to use Array API functionality/enhancements. The GFDLRegressor is not directly touched, as the array API modifications happen in base GFDL and invoked via super().fit() and super().predict(). From initial experiments, using accelerator devices with PyTorch or CuPy noticeably speeds up fit calls.

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

Test functions in test_regression.py are written to evaluate behavior with namespaces, devices, and dtypes that are accessed via sklearn.utils._array_api.yield_namespace_device_dtype_combinations(). We have tested these changes on NumPy, Array API Strict, PyTorch, and CuPy with CPU, MPS, and CUDA devices and float32 and float64 dtypes.

To test, you would use an incantation like:

python -m pip install -e .
python -m pip install array-api-strict array-api-compat torch cupy-cuda13x
SCIPY_ARRAY_API=1 PYTORCH_ENABLE_MPS_FALLBACK=1 python -m pytest

The changes use private utilities from sklearn.

---

Checklist:
[x] Benchmark speed against NumPy
[x] How important is solver=svd?
- It's more robust in matching values between NumPy and Torch backends. Runtime differences appear to be negligible. Documentation says SVD is more stable, which could be good when inverting very large p > n matrices.
[x] Are we able at all to check coefficients and predictions for PyTorch with bounded weight schemes?
- Yes, but we have to gain control of value checking in check_array_api_input_and_values from scikit-learn via a custom test function.
[x] Is there a more elegant solution than the private utility functions?
- move_to, _max_precision_float_dtype
[x] Is there an sklearn decorator we could use in the test suite?
- You add tags to the estimator class and use array_api_dispatch
[x] Would some of the test functions already be captured by an sklearn decorator?
- Yes, check(estimator) with array_api_dispatch already does a lot input and output validation
[x] Do the activation and weight generation functions need to be in xp?
- Weights come from NumPy and are then converted
- Activations should be in xp, except for scipy.special which already handles that
[x] Can we lower the atol in the test_torch_matches_numpy function?
- Removed this function for check_array_api_input_and_values

---

SI disclosure: Provided advice and small code snippets for which sklearn utilities we could leverage for model.py

[sdtemple]

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.

[tylerjereddy]

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.

[eiviani-lanl]

There is also a related issue in the regularized path. With solver="auto", Ridge selects Cholesky for NumPy input but SVD for Array API input, which makes equivalence testing difficult.

See my modified tests (based heavily on https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/linear_model/tests/test_logistic.py#L2858) below the fold:

Details

@pytest.mark.parametrize(
    "reg_alpha, n_features, hidden_layer_sizes",
    [
        (1e-1, 40, (100,)),
        (1e-1, 400, (100,)),
        (2, 40, (100,)),
        (2, 400, (100,)),
        (None, 40, (100,)),
        (None, 400, (100,)),
        (1e-1, 40, (100, 100,)),
        (1e-1, 400, (100, 100,)),
        (2, 40, (100, 100,)),
        (2, 400, (100, 100,)),
        (None, 40, (100, 100,)),
        (None, 400, (100, 100,)),
    ]
)
@pytest.mark.parametrize(
    "array_namespace, device_name, dtype_name",
    yield_namespace_device_dtype_combinations(),
)
def test_gfdl_array_api_compliance(
    reg_alpha, 
    n_features, 
    hidden_layer_sizes,
    array_namespace,
    device_name,
    dtype_name,
):
    xp, device = _array_api_for_tests(array_namespace, device_name, dtype_name)
    X_np, y_np = make_regression(
            n_samples=10_000,
            n_features=n_features,
            n_informative=int(n_features / 10),
            n_targets=1,
            random_state=0,
        )
    
    y_np /= np.std(y_np) # normalize the regression targets
    
    X_np = X_np.astype(dtype_name, copy=False)
    y_np = y_np.astype(dtype_name, copy=False)
    
    X_xp = xp.asarray(X_np, device=device)
    y_xp_or_np = np.asarray(y_np, copy=True)

    params = dict(reg_alpha = reg_alpha, hidden_layer_sizes = hidden_layer_sizes, seed=0)

    with config_context(array_api_dispatch=False):
        gfdl_np = GFDLRegressor(**params).fit(
            X_np, y_np
        )
        prediction_np = gfdl_np.predict(X_np)

    atol = _atol_for_type(dtype_name)*4
    rtol = 1e-4 if dtype_name == "float32" else 1e-8

    with config_context(array_api_dispatch=True):
        gfdl_xp = GFDLRegressor(**params).fit(
            X_xp, y_xp_or_np
        )
        attr_xp = getattr(gfdl_xp, "coeff_")
        attr_np = getattr(gfdl_np, "coeff_")
        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)

        prediction_xp = gfdl_xp.predict(X_xp)
        prediction_xp = move_to(prediction_xp, xp=np, device="cpu")
        assert_allclose(prediction_xp, prediction_np, rtol=rtol, atol=atol)

UserWarning: Using Array API dispatch to namespace array_api_strict with `solver='auto'` will result in using the solver 'svd'. The results may differ from those when using a Numpy array, because in that case the preferred solver would be cholesky. Set `solver='svd'` to suppress this warning.

[sdtemple]

Emma and I talked about this in-person. We are both comfortable setting the solver as svd. On their page, sklearn explicitly says that svd is the most stable solver, especially for nearly singular matrices. I think svd may be the default (what auto would get you) for many namespaces, but numpy uses cholesky. We have had some instances of concern for very large singular matrices, so stability could be a good call. However, the tradeoff is speed.

https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.Ridge.html

In scope, I believe, is to do a small speed benchmark on svd versus cholesky.

[sdtemple]

I did do a small speed test for numpy (with numbers similar to #129 (comment))

The speed difference is negligible.

I would thus propose we accept the change to use svd:

• No major speed difference
• More stable, especially for different backends doing matrix inversion on large matrices

[sdtemple]

And we convert integer arrays to float arrays elsewhere in the code.

This comment could probably be reviewed and resolved?

[sdtemple]

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]

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."

[tylerjereddy]

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.

[sdtemple]

I have been able to lean on the existing sklearn utilities to achieve what the private utilities I wrote in the initial commit did. Following these changes, I do not even use array_api_compat anymore, as the utility functions in sklearn depend on their vendoring sklearn/externals/array_api_compat.

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

I am now conflicted between using the sklearn utilities versus using the non-vendored array_api_compat. Aren't these ideas in conflict with each other?

[tylerjereddy]

There are indeed some messy problems here, since depending on private sklearn utilities (usually "private" means names/modules/namespaces starting with underscores, like _name, not made for public consumption) that may disappear at any time is an inherently bad idea because of the risk of breaking.

For now, this WIP PR is a template that Emma might use to transplant over to her PR and help us get merged in to sklearn proper, and there that risk goes away because they can't remove the private utilities without also fixing the accompanying sources that depend on them in sklearn proper.

To move forward here, this may be one of the last decisions we make on the GFDL side before merging. If Emma's PR gets merged first, then the pressure here to leverage sklearn utilities drops quite a bit, but I didn't want two people pursuing different approaches to the same goal, at least not yet, because then the synergy of effort is lost. If her PR does get merged first though, then we can probably do less painful things here. For example, we'd have the option of using array_api_compact directly in our library/source code so it doesn't break for consumers, but we could still keep the vulnerability of using sklearn private utils in testing code--since user/consumer breaks that only affect the testing code are far less serious.

But maybe for now make this PR as useful as possible for Emma's transplantation of array API capabilities is the right target to achieve the goal of getting the sklearn PR merged, which would be huge.

[tylerjereddy]

We could also try "dumping" the necessary private utilities into array-api-extra as noted at #129 (comment) two weeks ago. The community might appreciate that since others will have this problem too, but would need to determine how much work that would be. Then we wouldn't have any "private" deps and yet we'd still use an approach that is basically harmonious with sklearn/what Emma needs.

[sdtemple]

So these are the core private utilities that we use:

from sklearn.utils._array_api import (
    _max_precision_float_dtype,
    check_same_namespace,
    get_namespace_and_device,
    move_to,
)

I suspect these would be generally useful in something like array-api-extra.

[sdtemple]

@eiviani-lanl In that case, I think the current status of this PR that exclusively uses the sklearn private utilities

from sklearn.utils._array_api import (
    _max_precision_float_dtype,
    check_same_namespace,
    get_namespace_and_device,
    move_to,
)

and the testing utilities

from sklearn.utils._testing import (
    skip_if_array_api_compat_not_configured,
)
from sklearn.utils.estimator_checks import (
    check_array_api_input_and_values, ### THIS ONE IS THE MOST IMPORTANT
    parametrize_with_checks,
)

would be the best approach to take for your sklearn PR. As Tyler mentioned, they'd have to fix many downstream issues if removing those sklearn.utils._array_api functions, and those are used all over the place.

[tylerjereddy]

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.

[sdtemple]

I think this can be resolved. We are using sklearn private utilities now.

[tylerjereddy]

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]

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.

[sdtemple]

I am now using check_same_namespace

[sdtemple]

I think this old comment can be resolved.

[tylerjereddy]

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

[sdtemple]

This custom test has been removed in favor of sklearn utilities

[sdtemple]

I think this old comment can be resolved.

[tylerjereddy]

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

[sdtemple]

I removed this test in favor of sklearn tests.

[sdtemple]

I think this old comment can be resolved. That custom test was removed, and we use xp now instead of custom imports of namespaces.

[tylerjereddy]

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.

[sdtemple]

The toggle I am working with is

SCIPY_ARRAY_API=1 python -m pytest

This seems to be the way that sklearn performs some of their testing:

@pytest.mark.skipif(
    os.environ.get("SCIPY_ARRAY_API") != "1", reason="SCIPY_ARRAY_API not set to 1."
)

I am implementing that decorator in the array API-related tests in alignment with the sklearn approach.

[sdtemple]

I think this old comment can be resolved. We removed

os.environ["SCIPY_ARRAY_API"] = "1"

from code in the tests.

[tylerjereddy]

Note also that many of the tensions here with respect to gracefully extracting the xp-related testing utilities from sklearn/SciPy were discussed this morning in the Python Array API standard meeting (minutes should be public: https://hackmd.io/zn5bvdZTQIeJmb3RW1B-8g#Meeting-minutes-17-September-2026).

In particular, here: https://hackmd.io/zn5bvdZTQIeJmb3RW1B-8g#Extracting-common-testing-infrastructure-Lucas

array-api-extra currently has https://data-apis.org/array-api-extra/api-testing.html and https://data-apis.org/array-api-extra/api-lazy.html, but none of the fancier decorators for parametrization/yielding/skipping yet (per discussion with @lucascolley).

[eiviani-lanl]

@sdtemple In my own testing it's occurred to me that we'll need to make all activation functions namespace aware, i.e.

--- a/src/gfdl/activations.py
+++ b/src/gfdl/activations.py
@@ -20,7 +20,9 @@ def relu(z):
     numpy.ndarray
         The output array with only positive values.
     """
-    return np.maximum(0, z)
+    xp, _, device = get_namespace_and_device(z)
+    zero = xp.asarray(0, dtype=z.dtype, device=device)
+    return xp.maximum(z, zero)

[sdtemple]

Thank you for answering that item from the checklist items! I'll be sure to make that change.

I know we could also do array API support for weight generation, but those can be converted anyway by move_to. So, for that, it'd be more about speed testing.

[sdtemple]

Thank you for answering that item from the checklist items! I'll be sure to make that change.

I know we could also do array API support for weight generation, but those can be converted anyway by move_to. So, for that, it'd be more about speed testing.

Ignore this comment. It looks like the community discussed array API standards for random number generation at length and couldn't come to a consensus on implementation because of irreconcilable differences between packages like NumPy versus those like JAX.
data-apis/array-api#431

[tylerjereddy]

Yeah there is https://github.com/NeilGirdhar/array-api-rng, probably not going to be solved anytime soon. Best bet is, as usual, to look what sklearn/SciPy do and mimic that.

[tylerjereddy]

@sdtemple had asked in person how we can check which SciPy functions have array API support. Taking one of the ones we use for activation as an example, scipy.special.log_softmax, you can search for it in the SciPy devdocs (see here: https://scipy.github.io/devdocs/reference/generated/scipy.special.log_softmax.html#scipy.special.log_softmax).

Many functions like that one will have an array API standard support table/summary, which is quite nice. I'll reproduce it below. An overview of array API support coverage across the entire library is available in a larger table at https://scipy.github.io/devdocs/dev/api-dev/array_api.html#api-coverage.

[sdtemple]

@tylerjereddy This PR is ready for review. I don't know if it would make sense to:

• Change title to Array API for GFDLRegressor, since we use yield_namespace_device_dtype_combinations in the test and have found locally that it works for cuda, mps, and cupy
• For branch cleanliness, squash some of the commits, especially the early ones before sklearn utilities were being used

[tylerjereddy]

We can always squash the commits before merging, so not too worried about the commit history. Cleaning up the PR title and original comment/description for clarity can be helpful for the reviewer if they are out of date.

Having CI passing is also nice if something is genuinely ready to go. That depends on some other things like gh-117 and gh-98 though.

Depending on private sklearn utilities is the other obvious blocker I know about before diving back in. Even if the code were perfect, that would be a point of hesitation for merging. It would probably take quite some amount of time to get everything we need in array-api-extra. Another option would be some form of vendoring the sklearn utilities that we use--basically maintaining our own copy of that "private" infrastructure so we don't depend on the private sklearn version, and then it can't just randomly break when they make changes/do a release. If that made sense to do (it might), one question might be whether that code should be vendored into our source tree here directly, or as a git submodule that pulls in another repo with that code, or something like that.

[sdtemple]

Having CI passing is also nice if something is genuinely ready to go. That depends on some other things like #117 and #98 though.

Yes, those are the unrelated issues that cause the CI failures. From my repeated testing on Mac OS and the Linux box, #98 is not a failure locally.

Depending on private sklearn utilities is the other obvious blocker I know about before diving back in.

What would be the course of action then? Hold off for a bit? The code right now with sklearn private utilities works for research on speed testing. Initially, I don't like the idea of vendoring; it would be nice to see what progress could happen on this data-apis/array-api-extra#1064

[tylerjereddy]

#98 is not a failure locally.

It depends on the status of the servers/ecosystem at any given time.

What would be the course of action then? Hold off for a bit?

Well, if we don't vendor I bet we'll be waiting weeks-months for array-api-extra to merge our stuff in. Maybe that's pessimistic, but I dunno, probably would take quite a while to do it correctly. If there's no rush to merge this and you just want review comments/feedback, that's fine I suppose. There may be other reasons to want to merge though--like "certifying" that we are confident the results are correct with other backends, having them included in the next release for project progress highlights, and so on.

[tylerjereddy]

Also, I'm seeing new CI failures here that don't have issues open or explanations for them yet--try to clear those up before pinging for review.

The most important CI job here, "Array API CI," shows no tests ran in 0.03s--that doesn't look like it is ready for review yet.

You also have unrelated CI failures that you should try to help the team triage. I just opened gh-145 to describe some of those new failures that you can see in the Python 3.13 macos-latest job here.

It would also be good to have another team member confirm that they are fully happy with all the changes here before pinging me as noted last month at #129 (review). I'll try to get you some feedback pretty soon either way though, but for sustainability that's probably useful to do.