WIP, ENH: exact solve with array backends

[tylerjereddy]

• Related to ENH: Ridge supports array API standard, we could as well? #68, but for the opposite case of the exact solve.

• This is a (very) early draft/prototype of supporting the Python array API standard (https://data-apis.org/array-api/latest/) for GFDLClassifier when the exact solve is used. Some of the sklearn API conformance tests will currently fail, but all numerical correctness tests with the NumPy backend should continue to pass.

• The basic idea is to take an early peek at our potential for speedups when using JIT compiled/GPU backend array types, since i.e., hyperopt can be slow on CPU alone, and also because we'd just like to deliver the most performant estimators that we can, which also helps us conduct numerical experiments for the mathematicians with faster turnaround times.

Using a benchmark script similar to the one below the fold, here are some sample results showing some promising-looking speedups on this branch vs. plain NumPy. I'd probably prefer to hand this effort off to someone else, but does seem worth pursuing.

Details

import time
import argparse


from tqdm import tqdm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_classification
from sklearn import config_context
from sklearn.utils.validation import check_is_fitted
import jax.numpy as jnp
import torch


from gfdl.model import GFDLClassifier


def perform_benchmarks():
    results = {}

    sample_sizes = [1_000, 2_000, 3_000, 10_000]
    fig, ax = plt.subplots()
    for backend in tqdm(["numpy", "jax", "torch"], desc="backends"):
        average_times = []
        std_times = []
        for n_samples in tqdm(sample_sizes, desc="sample sizes"):
            X, y = make_classification(n_samples=n_samples, n_features=2000, random_state=42)

            if backend == "jax":
                X = jnp.array(X)
                y = jnp.array(y)
            elif backend == "torch":
                X = torch.asarray(X)
                y = torch.asarray(y)

            trial_times = []
            for num_trials in tqdm(range(3), desc="replicate"):
                with config_context(array_api_dispatch=True):
                    clf = GFDLClassifier()
                    start = time.perf_counter()
                    clf.fit(X, y)
                    # we check that the estimator is actually fitted, as a precaution
                    # against backends that may not be sync'd/blocking
                    check_is_fitted(clf)
                    end = time.perf_counter()
                    trial_times.append(end - start)
            average_times.append(np.mean(trial_times))
            std_times.append(np.std(trial_times))
        ax.errorbar(sample_sizes,
                    average_times,
                    yerr=std_times,
                    fmt="o",
                    linestyle="-",
                    capsize=5,
                    label=backend)
    ax.set_ylabel("Average Time (s)")
    ax.set_xlabel("# Records")
    ax.legend()
    ax.set_title("Array API support for GFDLClassifier with exact solve")
    fig.savefig("bench.png", dpi=300)


if __name__ == "__main__":
    perform_benchmarks()

On M3 Max laptop (using CPU):

On an x86_64 Linux box (also using i9-13900K CPU only for now):

initial TODO:

• add more mature array API support, including for GPU/devices (currently fails on this branch when I try to use a GPU backend for torch)
• add proper testing for array API support (see i.e., array-api-strict and the array API testing harnesses in sklearn and SciPy, etc.)
• sort out how to get our regular sklearn conformance tests passing
• wait for SciPy 2.0.0 release in early 2027 to turn this on by default? (but perhaps allow it behind a flag until then?) -- see: https://discuss.scientific-python.org/t/making-array-types-support-public-scipy-2-0/2355
• check inference timings, the timings above are just for fit()

AI usage: I did not use AI to write the source code of this PR

[tylerjereddy]

not all sklearn operations support the array API standard yet, see: https://scikit-learn.org/stable/modules/array_api.html

[tylerjereddy]

so, I'm coercing back here, for now

[tylerjereddy]

If I try to force torch to do its work on a GPU on my Linux box with torch.set_default_device("cuda") in the benchmark script, I see an error here (below). Fixing that will cascade to other errors that need to be solved, so this will be a bit of a project (and require matching tests).

However, the reward could be substantial for large design matrices. Of course, we don't even have direct solve as an option in any of our current hyperopt runs (that's a separate internal matter at: https://lisdi-git.lanl.gov/gfdl_research/research_demo/-/work_items/6).

Traceback (most recent call last):
  File "/home/treddy/rough_work/RVFL/June_23_2026/probe.py", line 71, in <module>
    perform_benchmarks()
    ~~~~~~~~~~~~~~~~~~^^
  File "/home/treddy/rough_work/RVFL/June_23_2026/probe.py", line 48, in perform_benchmarks
    clf.fit(X, y)
    ~~~~~~~^^^^^^
  File "/home/treddy/github_projects/ascr_rvfl/src/gfdl/model.py", line 424, in fit
    Y = self.enc_.fit_transform(np.from_dlpack(Y).reshape(-1, 1))
                                ~~~~~~~~~~~~~~^^^
BufferError: Unsupported device in DLTensor.

[tylerjereddy]

The most recent commit allows me to run benchmark code (below the fold) with torch + 4090 GPU on this branch, and still only fails the sklearn API conformance tests, but no numerical tests with NumPy.

The sample results with GPU (also below) aren't stunning yet, probably because of two things:

1. need really large data sizes to justify GPU usage (we should probably try this on even larger data to find the crossover point where torch GPU > CPU)
2. the code is probably doing unnecessary device transfers/coercions/copies, that need to be cleaned up...

Still, this is a step toward prototyping viability on GPU backends.

Details

import time
import argparse


from tqdm import tqdm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_classification
from sklearn import config_context
from sklearn.utils.validation import check_is_fitted
import jax.numpy as jnp
import torch


from gfdl.model import GFDLClassifier


def perform_benchmarks():
    results = {}

    sample_sizes = [1_000, 2_000, 10_000, 20_000, 50_000, 100_000]
    fig, ax = plt.subplots()
    for backend in tqdm(["torch gpu", "torch cpu", "numpy"], desc="backends"):
        average_times = []
        std_times = []
        for n_samples in tqdm(sample_sizes, desc="sample sizes"):
            X, y = make_classification(n_samples=n_samples, n_features=2000, random_state=42)

            if backend == "jax":
                X = jnp.array(X)
                y = jnp.array(y)
            elif backend == "torch gpu":
                torch.set_default_device("cuda")
                X = torch.asarray(X)
                y = torch.asarray(y)
            elif backend == "torch cpu":
                torch.set_default_device("cpu")
                X = torch.asarray(X)
                y = torch.asarray(y)
            elif backend == "cupy":
                X = cp.asarray(X)
                y = cp.asarray(y)

            trial_times = []
            for num_trials in tqdm(range(3), desc="replicate"):
                with config_context(array_api_dispatch=True):
                    clf = GFDLClassifier()
                    start = time.perf_counter()
                    clf.fit(X, y)
                    # we check that the estimator is actually fitted, as a precaution
                    # against backends that may not be sync'd/blocking
                    check_is_fitted(clf)
                    end = time.perf_counter()
                    trial_times.append(end - start)
            average_times.append(np.mean(trial_times))
            std_times.append(np.std(trial_times))
        ax.errorbar(sample_sizes,
                    average_times,
                    yerr=std_times,
                    fmt="o",
                    linestyle="-",
                    capsize=5,
                    label=backend)
    ax.set_ylabel("Average Time (s)")
    ax.set_xlabel("# Records")
    ax.legend()
    ax.set_title("Array API support for GFDLClassifier fit() with exact solve")
    fig.savefig("bench.png", dpi=300)


if __name__ == "__main__":
    perform_benchmarks()

[tylerjereddy]

This shim is not correct and breaks device handling on JAX IIRC. It is a temporary hack to allow collecting GPU benchmarks with torch. This is also obviously inefficient with GPU-CPU-GPU transfer happening...

But I'm trying to nudge this forward a bit with initial prototyping... especially with potential need to justify GPU resources in our next IC proposal

[tylerjereddy]

I pushed up a short paragraph to the manuscript/Overleaf describing this very early capability:

it should probably be improved to qualify that this is only for the exact solve, etc.

[tylerjereddy]

Because of the limitation that our current in house hyperot runs are only using Ridge codepath (see internal ticket: https://lisdi-git.lanl.gov/gfdl_research/research_demo/-/work_items/6), and the above performance results are only for the pinv codepath, we have some interest in showing that we could potentially use the Python Array API standard (and GPUs) to gain a performance advantage with Ridge on i.e., HPC hardware with GPUs. If we could demonstrate this, it may be ammunition for our institutional computing proposals to request appropriate GPU resources, and provide us with faster iteration times for our hyperopt research explorations.

One challenge we face here is that the code in this draft/WIP PR obviously has some undesirable device transfers (host<->GPU/device data transfers) that will invariably cause performance loss. We need to either clean that up, and add appropriate array API-focused testing, or have a clear justification for large enough datasets that the GPU compensates for the cost of the device transfer (realistically, we may need a bit of both).

Lucy's recent Quansight/sklearn joint blogpost on Array API support in sklearn (https://labs.quansight.org/blog/array-api-scikit-learn-2026), which we've discussed in meetings over the past few months, indicates that Ridge has had array API support for awhile, but suggests that we need to adopt a specific solver for now:

Note that scikit-learn’s Ridge regressor currently only supports ‘svd’ solver. We selected this solver for initial implementation as it exclusively uses standard-compliant functions available across all backends and is the most stable solver. Support for the ‘cholesky’ solver is also underway

So, I took the feature branch associated with this PR, and used the following small diff for now, to enforce usage of that solver:

--- a/src/gfdl/model.py
+++ b/src/gfdl/model.py
@@ -109,7 +109,7 @@ class GFDL(BaseEstimator):
         if self.reg_alpha is None:
             self.coeff_ = xp.linalg.pinv(D, rtol=self.rtol) @ Y
         else:
-            ridge = Ridge(alpha=self.reg_alpha, fit_intercept=False)
+            ridge = Ridge(alpha=self.reg_alpha, fit_intercept=False, solver="svd")
             ridge.fit(D, Y)
             self.coeff_ = ridge.coef_.T
         return self

and then combined that with the benchmark script below the fold, which ran in about 9 mintues on gp160.lanl.gov, which has an i9-13900K CPU and 4090 GPU.

Details

import time
import argparse


from tqdm import tqdm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_classification
from sklearn import config_context
from sklearn.utils.validation import check_is_fitted
import jax.numpy as jnp
import torch


from gfdl.model import GFDLClassifier


def perform_benchmarks():
    results = {}

    sample_sizes = [10_000, 20_000, 30_000, 40_000, 60_000, 100_000, 200_000]
    fig, ax = plt.subplots()
    for backend in tqdm(["numpy", "torch", "torch (gpu)"], desc="backends"):
        average_times = []
        std_times = []
        for n_samples in tqdm(sample_sizes, desc="sample sizes"):
            X, y = make_classification(n_samples=n_samples, n_features=2000, random_state=42)

            if backend == "jax":
                X = jnp.array(X)
                y = jnp.array(y)
            elif backend == "torch":
                torch.set_default_device('cpu')
                X = torch.asarray(X)
                y = torch.asarray(y)
            elif backend == "torch (gpu)":
                torch.set_default_device('cuda')
                X = torch.asarray(X)
                y = torch.asarray(y)

            trial_times = []
            for num_trials in tqdm(range(3), desc="replicate"):
                with config_context(array_api_dispatch=True):
                    clf = GFDLClassifier(reg_alpha=0.1)
                    start = time.perf_counter()
                    clf.fit(X, y)
                    # we check that the estimator is actually fitted, as a precaution
                    # against backends that may not be sync'd/blocking
                    check_is_fitted(clf)
                    end = time.perf_counter()
                    trial_times.append(end - start)
            average_times.append(np.mean(trial_times))
            std_times.append(np.std(trial_times))
        ax.errorbar(sample_sizes,
                    average_times,
                    yerr=std_times,
                    fmt="o",
                    linestyle="-",
                    capsize=5,
                    label=backend)
    ax.set_ylabel("Average Time for estimator fit() (s)")
    ax.set_xlabel("# Records")
    ax.legend()
    ax.set_title("Array API support for GFDLClassifier with Ridge")
    fig.savefig("bench.png", dpi=300)


if __name__ == "__main__":
    perform_benchmarks()

We can see from the results below that torch on the CPU already provides a nice speedup over native NumPy on CPU. The GPU does outperform the CPU eventually, but requires quite a large design matrix, as expected for the reasons outlined above. However, we're still not getting huge speedups on GPU--I suspect more care is needed here in terms of the adaptation of the array API standard in a manner that does not leverage device transfers, etc. Still, a step in the direction of demonstrating that hyperopt runs will likely benefit from the array API standard, though we should carefully check the svd solver limitation re: gracefully falling back to CPU/NumPy with other solvers (probably it would just error out...). Interestingly, it appears that limiting Ridge to svd in this way doesn't currently add more test failures to the suite on this branch on x86_64 Linux.

Probably someone with more bandwidth than I should champion our adoption of the array API standard, which is likely a decent chunk of work re: building up the testsuite, JAX, torch, CuPy support, etc.

[tylerjereddy]

Also, one important thing I didn't mention yet--if we're being held back by codepaths in sklearn that do not yet support the array API standard, it would of course benefit everyone in the community for us to try adding that support.

[tylerjereddy]

the balanced E-MNIST dataset we use for hyperopt does have > 100,000 records, so we may already be in the relevant regime without needing an excuse.

[tylerjereddy]

This morning one of the sklearn developers, Olivier, expressed an interest in the status of our support for the Python array API standard (see: scikit-learn/scikit-learn#33215 (comment)).

I'll need to think about who has the bandwidth and interest to drive this forward in the near term, if it helps get that work across the line.