From 1661dea7c8103155275c11abe9ec9378ef71f810 Mon Sep 17 00:00:00 2001 From: Seth Temple Date: Sat, 12 Sep 2026 12:48:45 -0600 Subject: [PATCH 1/2] WIP: array api standard for torch cpu - 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 --- pyproject.toml | 2 + src/gfdl/model.py | 17 ++++-- src/gfdl/tests/test_regression.py | 99 +++++++++++++++++++++++++++++++ 3 files changed, 114 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d566d9c..6de9dda 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,8 @@ classifiers = [ dependencies = ["numpy>=2.0.0", "scikit-learn>=1.5.0", "scipy>=1.13.0", + "array-api-compat", + "torch>=2.12.0,<=2.14.0", "packaging>=24.0"] [project.urls] diff --git a/src/gfdl/model.py b/src/gfdl/model.py index 56dd28a..9b22bbd 100644 --- a/src/gfdl/model.py +++ b/src/gfdl/model.py @@ -2,6 +2,7 @@ Estimators for gradient free deep learning. """ +import array_api_compat import numpy as np import scipy from scipy.special import logsumexp @@ -44,12 +45,14 @@ def __init__( self.rtol = rtol def fit(self, X, Y): + xp = array_api_compat.get_namespace(X, Y) + # Assumption : X, Y have been pre-processed. # X shape: (n_samples, n_features) # Y shape: (n_samples, n_classes-1) if self.reg_alpha is not None and self.reg_alpha < 0.0: raise ValueError("Negative reg_alpha. Expected range : None or [0.0, inf).") - hidden_layer_sizes = np.asarray(self.hidden_layer_sizes) + hidden_layer_sizes = xp.asarray(self.hidden_layer_sizes) if hidden_layer_sizes.min() < 1: raise ValueError("hidden_layer_sizes must be > 0, " f"got {hidden_layer_sizes}") @@ -87,6 +90,8 @@ def fit(self, X, Y): Hs = [] H_prev = X for w, b in zip(self.W_, self.b_, strict=False): + w = xp.asarray(w, device=H_prev.device,) + b = xp.asarray(b, device=H_prev.device,) Z = H_prev @ w.T + b # (n_samples, n_hidden) H_prev = self._activation_fn(Z) Hs.append(H_prev) @@ -95,7 +100,7 @@ def fit(self, X, Y): # or (n_samples, sum_hidden) if self.direct_links: Hs.append(X) - D = np.hstack(Hs) + D = xp.concat(Hs, axis=1) # beta shape: (sum_hidden+n_features, n_classes-1) # or (sum_hidden, n_classes-1) @@ -103,7 +108,7 @@ def fit(self, X, Y): # If reg_alpha is None, use direct solve using # 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 else: ridge = Ridge(alpha=self.reg_alpha, fit_intercept=False) ridge.fit(D, Y) @@ -229,16 +234,20 @@ def partial_fit(self, X, y): def predict(self, X): check_is_fitted(self) + xp = array_api_compat.get_namespace(X) + Hs = [] H_prev = X for W, b in zip(self.W_, self.b_, strict=False): + W = xp.asarray(W, device=H_prev.device) + b = xp.asarray(b, device=H_prev.device) Z = H_prev @ W.T + b # (n, m) H_prev = self._activation_fn(Z) Hs.append(H_prev) if self.direct_links: Hs.append(X) - D = np.hstack(Hs) + D = xp.concat(Hs, axis=1) out = D @ self.coeff_ return out diff --git a/src/gfdl/tests/test_regression.py b/src/gfdl/tests/test_regression.py index c0cab55..5a759fc 100644 --- a/src/gfdl/tests/test_regression.py +++ b/src/gfdl/tests/test_regression.py @@ -1,5 +1,11 @@ +import os + +import array_api_compat import numpy as np import pytest +import torch +from numpy.testing import assert_allclose +from sklearn import config_context from sklearn.base import clone from sklearn.datasets import fetch_openml, make_regression from sklearn.metrics import r2_score @@ -257,3 +263,96 @@ def test_preserve_class_inputs(): for k, v in actual.items(): assert v == expected[k] assert isinstance(v, type(expected[k])) + + +@pytest.mark.parametrize( + "estimator, reg_alpha, n_features", + [ + (GFDLRegressor, 1e-1, 40), + (GFDLRegressor, 1e-1, 400), + (GFDLRegressor, 2, 40), + (GFDLRegressor, 2, 400), + (GFDLRegressor, None, 40), + (GFDLRegressor, None, 400), + ] +) +def test_torch_cpu_matches_numpy(estimator, + reg_alpha, + n_features, + ): + # numpy and torch accuracies should be close up to atol + os.environ["SCIPY_ARRAY_API"] = "1" + report = r2_score + rng = np.random.default_rng(seed=42) + acc_np_s = [] + acc_torch_s = [] + for _ in range(10): + random_state = rng.integers(low=0, high=100, size=1)[0] + 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=random_state, + ) + X_torch = torch.asarray(X_np, device="cpu",) + y_torch = torch.asarray(y_np, device="cpu",) + with config_context(array_api_dispatch=True): + model = estimator(reg_alpha=reg_alpha) + model.fit(X_np, y_np) + y_pred = model.predict(X_torch) + acc_torch = report( + y_torch, + y_pred, + ) + acc_torch_s.append(acc_torch) + # test against numpy which is current usage + with config_context(array_api_dispatch=True): + model = estimator(reg_alpha=reg_alpha) + model.fit(X_np, y_np) + y_pred = model.predict(X_np) + acc_np = report( + y_np, + y_pred, + ) + acc_np_s.append(acc_np) + assert_allclose(acc_torch_s, acc_np_s, atol=1e-3) + + +@pytest.mark.parametrize( + "estimator, namespace, device, X_dtype, y_dtype", + [ + (GFDLRegressor, np, "cpu", np.float64, np.float64), + (GFDLRegressor, torch, "cpu", torch.float64, torch.float64), + ] +) +def test_predict_same_namespace(estimator, + namespace, + device, + X_dtype, + y_dtype, + ): + # ValueError thrown .predict has + # different namespace from X and y + os.environ["SCIPY_ARRAY_API"] = "1" + + X, y = make_regression( + n_samples=1_000, + n_features=10, + n_informative=4, + n_targets=1, + random_state=42, + ) + X = namespace.asarray(X, dtype=X_dtype, device=device) + y = namespace.asarray(y, dtype=y_dtype, device=device) + xp = array_api_compat.get_namespace(X, y) + with config_context(array_api_dispatch=True): + model = estimator() + model.fit(X, y) + y_pred = model.predict(X) + + xp_y_pred = array_api_compat.get_namespace(y_pred) + if xp_y_pred != xp: + raise ValueError( + ".predict output and X, y are not in the same namespace" + ) From 6d9beacc77c9facf290c44a2cc39128dafb8faa7 Mon Sep 17 00:00:00 2001 From: Seth Temple Date: Sun, 13 Sep 2026 22:32:02 -0600 Subject: [PATCH 2/2] WIP: ensure float, check and convert arrays - 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 --- src/gfdl/model.py | 75 +++++++++++-- src/gfdl/tests/test_regression.py | 172 ++++++++++++++++++++++++------ 2 files changed, 204 insertions(+), 43 deletions(-) diff --git a/src/gfdl/model.py b/src/gfdl/model.py index 9b22bbd..6e6358f 100644 --- a/src/gfdl/model.py +++ b/src/gfdl/model.py @@ -45,7 +45,9 @@ def __init__( self.rtol = rtol def fit(self, X, Y): - xp = array_api_compat.get_namespace(X, Y) + xp = array_api_compat.get_namespace(X) + X = _ensure_float_X(xp, X) + Y = _check_and_convert_array(X, Y) # Assumption : X, Y have been pre-processed. # X shape: (n_samples, n_features) @@ -68,30 +70,40 @@ def fit(self, X, Y): rng = self.get_generator(self.seed) self.W_.append( + _check_and_convert_array( + X, self._weight_mode( self._N, hidden_layer_sizes[0], rng=self.get_generator(self.seed) - ) + ), + ) ) self.b_.append( + _check_and_convert_array( + X, self._weight_mode(1, hidden_layer_sizes[0], rng=rng) - .reshape(-1) + .reshape(-1), + ) ) for i, layer in enumerate(hidden_layer_sizes[1:]): # (n_hidden, n_features) self.W_.append( - self._weight_mode(hidden_layer_sizes[i], layer, rng=rng,) + _check_and_convert_array( + X, + self._weight_mode(hidden_layer_sizes[i], layer, rng=rng,), + ) ) # (n_hidden,) self.b_.append( - self._weight_mode(1, layer, rng=rng,).reshape(-1) + _check_and_convert_array( + X, + self._weight_mode(1, layer, rng=rng,).reshape(-1), + ) ) # hypothesis space shape: (n_layers,) Hs = [] H_prev = X for w, b in zip(self.W_, self.b_, strict=False): - w = xp.asarray(w, device=H_prev.device,) - b = xp.asarray(b, device=H_prev.device,) Z = H_prev @ w.T + b # (n_samples, n_hidden) H_prev = self._activation_fn(Z) Hs.append(H_prev) @@ -239,8 +251,8 @@ def predict(self, X): Hs = [] H_prev = X for W, b in zip(self.W_, self.b_, strict=False): - W = xp.asarray(W, device=H_prev.device) - b = xp.asarray(b, device=H_prev.device) + W = _check_and_convert_array(X, W) + b = _check_and_convert_array(X, b) Z = H_prev @ W.T + b # (n, m) H_prev = self._activation_fn(Z) Hs.append(H_prev) @@ -248,7 +260,10 @@ def predict(self, X): if self.direct_links: Hs.append(X) D = xp.concat(Hs, axis=1) - out = D @ self.coeff_ + out = D @ _check_and_convert_array( + X, + self.coeff_ + ) return out @@ -1275,3 +1290,43 @@ def predict(self, X): check_is_fitted(self) X = validate_data(self, X, reset=False) return super().predict(X) + + +def _check_and_convert_array(X1, X2): + """Convert second array to namespace, device, dtype of first if not already""" + xp = array_api_compat.get_namespace(X1) + kwargs = {} + + if hasattr(X1, "dtype") and getattr(X2, "dtype", None) != X1.dtype: + kwargs["dtype"] = X1.dtype + + x1_device = getattr(X1, "device", None) + x2_device = getattr(X2, "device", None) + + if x1_device is not None and x1_device != x2_device: + kwargs["device"] = x1_device + + if kwargs: + return xp.asarray(X2, **kwargs) + + return X2 + + +def _ensure_float_X(xp, X): + """Make design matrix floating point numbers""" + if xp.isdtype(X.dtype, "real floating"): + return X + + device = getattr(X, "device", None) + device_str = str(device).lower() + + if "cuda" in device_str or "gpu" in device_str or "tpu" in device_str: + dtype = xp.float32 + else: + dtype = xp.float64 + + kwargs = {"dtype": dtype} + if device is not None: + kwargs["device"] = device + + return xp.asarray(X, **kwargs) diff --git a/src/gfdl/tests/test_regression.py b/src/gfdl/tests/test_regression.py index 5a759fc..a7172e0 100644 --- a/src/gfdl/tests/test_regression.py +++ b/src/gfdl/tests/test_regression.py @@ -266,22 +266,29 @@ def test_preserve_class_inputs(): @pytest.mark.parametrize( - "estimator, reg_alpha, n_features", + "reg_alpha, n_features, hidden_layer_sizes", [ - (GFDLRegressor, 1e-1, 40), - (GFDLRegressor, 1e-1, 400), - (GFDLRegressor, 2, 40), - (GFDLRegressor, 2, 400), - (GFDLRegressor, None, 40), - (GFDLRegressor, None, 400), + (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,)), ] ) -def test_torch_cpu_matches_numpy(estimator, - reg_alpha, - n_features, - ): - # numpy and torch accuracies should be close up to atol +def test_torch_matches_numpy(reg_alpha, + n_features, + hidden_layer_sizes, + ): + """NumPy and Torch predictions should be close up to atol""" os.environ["SCIPY_ARRAY_API"] = "1" + estimator = GFDLRegressor report = r2_score rng = np.random.default_rng(seed=42) acc_np_s = [] @@ -298,7 +305,10 @@ def test_torch_cpu_matches_numpy(estimator, X_torch = torch.asarray(X_np, device="cpu",) y_torch = torch.asarray(y_np, device="cpu",) with config_context(array_api_dispatch=True): - model = estimator(reg_alpha=reg_alpha) + model = estimator(reg_alpha=reg_alpha, + hidden_layer_sizes=hidden_layer_sizes, + seed=random_state, + ) model.fit(X_np, y_np) y_pred = model.predict(X_torch) acc_torch = report( @@ -308,7 +318,10 @@ def test_torch_cpu_matches_numpy(estimator, acc_torch_s.append(acc_torch) # test against numpy which is current usage with config_context(array_api_dispatch=True): - model = estimator(reg_alpha=reg_alpha) + model = estimator(reg_alpha=reg_alpha, + hidden_layer_sizes=hidden_layer_sizes, + seed=random_state, + ) model.fit(X_np, y_np) y_pred = model.predict(X_np) acc_np = report( @@ -320,21 +333,24 @@ def test_torch_cpu_matches_numpy(estimator, @pytest.mark.parametrize( - "estimator, namespace, device, X_dtype, y_dtype", + "namespace, fit_dtype, predict_dtype", [ - (GFDLRegressor, np, "cpu", np.float64, np.float64), - (GFDLRegressor, torch, "cpu", torch.float64, torch.float64), + (np, np.float64, np.float32), + (torch, torch.float64, torch.float32), + (np, np.float32, np.float64), + (torch, torch.float32, torch.float64), ] ) -def test_predict_same_namespace(estimator, - namespace, - device, - X_dtype, - y_dtype, - ): +def test_predictor_context(namespace, + fit_dtype, + predict_dtype, + ): # ValueError thrown .predict has - # different namespace from X and y + # different namespace, dtype + # from its arguments os.environ["SCIPY_ARRAY_API"] = "1" + estimator = GFDLRegressor + device = "cpu" X, y = make_regression( n_samples=1_000, @@ -343,16 +359,106 @@ def test_predict_same_namespace(estimator, n_targets=1, random_state=42, ) - X = namespace.asarray(X, dtype=X_dtype, device=device) - y = namespace.asarray(y, dtype=y_dtype, device=device) - xp = array_api_compat.get_namespace(X, y) + X1 = namespace.asarray(X, dtype=fit_dtype, device=device) + X2 = namespace.asarray(X, dtype=predict_dtype, device=device) + y = namespace.asarray(y, dtype=fit_dtype, device=device) with config_context(array_api_dispatch=True): - model = estimator() - model.fit(X, y) - y_pred = model.predict(X) + y_pred = estimator().fit(X1, y).predict(X2) - xp_y_pred = array_api_compat.get_namespace(y_pred) - if xp_y_pred != xp: + xp_ypred = array_api_compat.get_namespace(y_pred) + xp_X2 = array_api_compat.get_namespace(X2) + if xp_ypred != xp_X2: raise ValueError( - ".predict output and X, y are not in the same namespace" + ".predict output and input are not in the same namespace" + ) + if predict_dtype != y_pred.dtype: + raise ValueError( + ".predict output and input are not the same dtype" + ) + + +@pytest.mark.parametrize( + "namespace, X_dtype, y_dtype", + [ + (np, np.float64, np.int64), + (np, np.int64, np.float64), + (torch, torch.float64, torch.int64), + ] +) +def test_int_array_api(namespace, + X_dtype, + y_dtype, + ): + """Integer arrays shall be handled gracefully""" + os.environ["SCIPY_ARRAY_API"] = "1" + estimator = GFDLRegressor + + X, y = make_regression( + n_samples=1_000, + n_features=10, + n_informative=4, + n_targets=1, + random_state=42, + ) + X = namespace.asarray(X, dtype=X_dtype, device="cpu") + y = namespace.asarray(y, dtype=y_dtype, device="cpu") + with config_context(array_api_dispatch=True): + estimator().fit(X, y) + + +@pytest.mark.parametrize( + "namespace, X_dtype, y_dtype", + [ + (np, np.float64, np.float64), + (np, np.float32, np.float32), + (torch, torch.float64, torch.float64), + (torch, torch.float32, torch.float32), + (torch, torch.float64, torch.float32), + ] +) +def test_fit_attr_context(namespace, + X_dtype, + y_dtype, + ): + """Fitted attributes shall be same as design matrix""" + os.environ["SCIPY_ARRAY_API"] = "1" + estimator = GFDLRegressor + device = "cpu" + + X, y = make_regression( + n_samples=1_000, + n_features=10, + n_informative=4, + n_targets=1, + random_state=42, + ) + X = namespace.asarray(X, dtype=X_dtype, device=device) + y = namespace.asarray(y, dtype=y_dtype, device=device) + with config_context(array_api_dispatch=True): + model = estimator().fit(X, y) + + dtypes = [] + devices = [] + spaces = [] + W_ = model.W_ + b_ = model.b_ + for w, b in zip(W_, b_, strict=False): + dtypes.extend((b.dtype, w.dtype)) + devices.extend((b.device, w.device)) + spaces.extend( + (array_api_compat.get_namespace(w), + array_api_compat.get_namespace(b),) ) + + coeff_ = model.coeff_ + dtypes.append(coeff_.dtype) + devices.append(coeff_.device) + spaces.append( + array_api_compat.get_namespace(coeff_) + ) + + assert all(dtype == X.dtype for dtype in dtypes) + assert all(device == X.device for device in devices) + + X_space = array_api_compat.get_namespace(X) + assert all(space == X_space for space in spaces)