From 1d78508eb6210902e902ca005b8fe1aa722daca0 Mon Sep 17 00:00:00 2001 From: mskumar Date: Fri, 25 Sep 2026 22:42:27 +0530 Subject: [PATCH 1/2] feat: preserve float dtypes and propagate CUDA devices end to end --- README.md | 14 +- RELEASES.md | 2 + benchmarks/compare_sklearn.py | 12 +- doc/index.rst | 1 + doc/user_guide/comparison.rst | 27 ++-- doc/user_guide/dtypes_devices.rst | 40 ++++++ tests/common/test_dtype_device.py | 132 +++++++++++++++++++ torml/base/_base.py | 2 - torml/cluster/_dbscan.py | 12 +- torml/cluster/_kmeans.py | 8 +- torml/covariance/_empirical.py | 16 ++- torml/cross_decomposition/_pls.py | 17 +-- torml/decomposition/_pca.py | 6 +- torml/discriminant_analysis/_lda.py | 20 +-- torml/ensemble/_bagging.py | 33 +++-- torml/ensemble/_forest.py | 24 ++-- torml/ensemble/_voting.py | 14 +- torml/feature_extraction/_dict_vectorizer.py | 5 +- torml/feature_selection/_univariate.py | 17 +-- torml/gaussian_process/_gpr.py | 7 +- torml/linear_model.py | 14 +- torml/manifold/_mds.py | 2 +- torml/mixture/_gmm.py | 21 ++- torml/model_selection/_split.py | 12 +- torml/multiclass/_ovr.py | 13 +- torml/multivariate/_multioutput.py | 13 +- torml/naive_bayes/_gaussian.py | 17 ++- torml/neighbors/_classification.py | 19 +-- torml/neighbors/_regression.py | 11 +- torml/pipelines/_union.py | 4 +- torml/preprocessing/_data.py | 12 +- torml/preprocessing/_encoders.py | 10 +- torml/random_projection/_gaussian.py | 8 +- torml/semi_supervised/_label_propagation.py | 10 +- torml/svm/_linear.py | 34 +++-- torml/tree/_tree.py | 45 ++++--- torml/utils/_mask.py | 3 +- torml/utils/_validation.py | 18 ++- 38 files changed, 471 insertions(+), 204 deletions(-) create mode 100644 doc/user_guide/dtypes_devices.rst create mode 100644 tests/common/test_dtype_device.py diff --git a/README.md b/README.md index 2a0ab27..32bb059 100644 --- a/README.md +++ b/README.md @@ -112,15 +112,17 @@ rerun with `uv run --with scikit-learn python benchmarks/compare_sklearn.py`): | task | metric | torml | scikit-learn | | --- | --- | --- | --- | -| linreg | R² | 0.999596 | 0.999596 | -| kmeans | inertia (lower better) | 14235.3 | 14271.8 | -| knn | accuracy | 0.9240 | 0.9240 | +| linreg | R² | 0.999656 | 0.999656 | +| kmeans | inertia (lower better) | 72645.4 | 72611.7 | +| knn | accuracy | 0.9476 | 0.9476 | | tree | accuracy | 1.0000 | 1.0000 | | scaler | max \|diff\| | 4.77e-07 | reference | -Correctness matches everywhere; timings are in the same class except tree -fitting (~2x slower — Cython vs Python loops). torml's real edges are -torch-native differentiable I/O, one runtime dependency, and exact +Correctness matches everywhere (10k samples per task); timings are in the +same class except tree fitting (~1.5x slower — Cython vs Python loops) and +10k kNN prediction, where sklearn's ball tree (~0.3s) beats torml's exact +brute force (~5s). torml's real edges are torch-native differentiable I/O +(including CUDA + float64 propagation), one runtime dependency, and exact statistics. Full honest write-up (including where sklearn wins): [user guide comparison](doc/user_guide/comparison.rst). diff --git a/RELEASES.md b/RELEASES.md index fcba395..8257930 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -3,6 +3,8 @@ ## [Unreleased] ### Added +- Added dtype/device propagation: `float32`/`float64` preserved end to end, CUDA inputs stay on CUDA; covered by `tests/common/test_dtype_device.py` (CUDA cases gated on availability). +- Added `doc/user_guide/dtypes_devices.rst` and scaled the sklearn comparison to 10k samples per task. ### Changed - Changed `torml.model_selection.GridSearchCV` docs: added `examples/model_selection/plot_grid_search.py`. diff --git a/benchmarks/compare_sklearn.py b/benchmarks/compare_sklearn.py index d063375..70f32b5 100644 --- a/benchmarks/compare_sklearn.py +++ b/benchmarks/compare_sklearn.py @@ -36,8 +36,8 @@ def bench_linreg(): from torml.metrics import r2_score torch.manual_seed(0) - features = torch.randn(5000, 20) - target = features @ torch.randn(20) + 0.1 * torch.randn(5000) + features = torch.randn(10000, 20) + target = features @ torch.randn(20) + 0.1 * torch.randn(10000) np_x, np_y = features.numpy(), target.numpy() tm, mine = timed(lambda: ToLin().fit(features, target)) ts, theirs = timed(lambda: SkLin().fit(np_x, np_y)) @@ -60,7 +60,7 @@ def bench_kmeans(): from torml.cluster import KMeans as ToKM torch.manual_seed(1) - data = torch.randn(2000, 10) + data = torch.randn(10000, 10) tm, mine = timed(lambda: ToKM(n_clusters=10, random_state=0, n_init=3).fit(data)) ts, theirs = timed( lambda: SkKM(n_clusters=10, random_state=0, n_init=3).fit(data.numpy()) @@ -81,7 +81,7 @@ def bench_knn(): from torml.neighbors import KNeighborsClassifier as ToKNN torch.manual_seed(2) - features = torch.randn(1000, 10) + features = torch.randn(10000, 10) labels = (features[:, 0] > 0).long() np_x, np_y = features.numpy(), labels.numpy() tm, mine = timed(lambda: ToKNN(5).fit(features, labels)) @@ -89,7 +89,7 @@ def bench_knn(): print(ROW.format("knn fit", "seconds (best of 3)", f"{tm:.4f}", f"{ts:.4f}")) tm, my_pred = timed(lambda: mine.predict(features)) ts, their_pred = timed(lambda: theirs.predict(np_x)) - print(ROW.format("knn predict-1000", "seconds", f"{tm:.4f}", f"{ts:.4f}")) + print(ROW.format("knn predict-10k", "seconds", f"{tm:.4f}", f"{ts:.4f}")) print( ROW.format( "knn", @@ -129,7 +129,7 @@ def bench_scaler(): from torml.preprocessing import StandardScaler as ToSS torch.manual_seed(0) - data = torch.randn(5000, 20) + data = torch.randn(10000, 20) tm, mine = timed(lambda: ToSS().fit_transform(data)) ts, theirs = timed(lambda: SkSS().fit_transform(data.numpy())) print(ROW.format("scaler", "seconds (best of 3)", f"{tm:.4f}", f"{ts:.4f}")) diff --git a/doc/index.rst b/doc/index.rst index 5734834..c574c89 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -18,6 +18,7 @@ Start with :doc:`quickstart`, learn the patterns in user_guide/unsupervised user_guide/preprocessing user_guide/model_selection + user_guide/dtypes_devices user_guide/comparison api diff --git a/doc/user_guide/comparison.rst b/doc/user_guide/comparison.rst index 36f8d8d..0fa55c9 100644 --- a/doc/user_guide/comparison.rst +++ b/doc/user_guide/comparison.rst @@ -9,7 +9,7 @@ Rerun it any time: uv run --with scikit-learn python benchmarks/compare_sklearn.py -Method: same data, same seeds, CPU, best of 3 timings +Method: same data, same seeds, CPU, best of 3 timings, 10k samples per task (Apple M-series, 2026-09-25, torch 2.14.0, scikit-learn 1.9.1). Results @@ -18,16 +18,16 @@ Results =================== ======================= ================== ================== task metric torml scikit-learn =================== ======================= ================== ================== -linreg fit seconds (best of 3) 0.0006 0.0009 -linreg R² 0.999596 0.999596 -kmeans fit seconds (best of 3) 0.0292 0.0278 -kmeans inertia (lower better) 14235.3 14271.8 -knn fit seconds (best of 3) 0.0000 0.0003 -knn predict-1000 seconds 0.0055 0.0057 -knn accuracy 0.9240 0.9240 -tree fit seconds (best of 3) 0.0010 0.0006 +linreg fit seconds (best of 3) 0.0012 0.0014 +linreg R² 0.999656 0.999656 +kmeans fit seconds (best of 3) 0.1631 0.1032 +kmeans inertia (lower better) 72645.4 72611.7 +knn fit seconds (best of 3) 0.0003 0.0018 +knn predict-10k seconds 5.1368 0.2828 +knn accuracy 0.9476 0.9476 +tree fit seconds (best of 3) 0.0083 0.0057 tree accuracy 1.0000 1.0000 -scaler seconds (best of 3) 0.0005 0.0005 +scaler seconds (best of 3) 0.0007 0.0008 scaler max |diff| vs sklearn 4.77e-07 reference torch-native backward() thru predict True n/a (numpy out) =================== ======================= ================== ================== @@ -35,7 +35,9 @@ torch-native backward() thru predict True n/a (numpy out) Reading the table honestly: correctness matches everywhere (identical R², accuracy, and scaler outputs; kmeans inertia differs only by random initialization). Timings are in the same class except tree fitting, where -scikit-learn's Cython is ~2x faster than torml's Python CART loops. +scikit-learn's Cython is ~1.5x faster than torml's Python CART loops, and +kNN prediction at 10k, where scikit-learn's ball tree (~0.3s) beats torml's +exact brute-force pairwise distances (~5s). Where torml is genuinely better ------------------------------- @@ -48,6 +50,9 @@ Where torml is genuinely better p-values are computed exactly, not via approximations. - **Readable from-scratch code** with strict input validation on every estimator; useful for teaching and auditing. +- **Dtypes and devices propagate.** ``float32``/``float64`` inputs keep + their precision end to end (integers promote to ``float32``), and CUDA + inputs stay on CUDA — see :doc:`dtypes_devices`. Where scikit-learn still wins ----------------------------- diff --git a/doc/user_guide/dtypes_devices.rst b/doc/user_guide/dtypes_devices.rst new file mode 100644 index 0000000..9b3ff43 --- /dev/null +++ b/doc/user_guide/dtypes_devices.rst @@ -0,0 +1,40 @@ +Dtypes and devices +==================== + +torml preserves your input dtype and device instead of silently converting +to ``float32``/CPU. + +Dtypes +------ + +- ``float32`` and ``float64`` tensors keep their precision end to end: + statistics (``mean_``, ``coef_``, ``components_``) and predictions come + back in the input dtype. +- Integer feature tensors are promoted to ``float32`` (means and variances + need fractions); integer *label* tensors stay integer. +- Lists default to ``float32``. Pass ``dtype=`` explicitly to + ``check_array`` to force a conversion. + +.. code-block:: python + + X64 = torch.randn(20, 3, dtype=torch.float64) + model = LinearRegression().fit(X64, y64) + model.coef_.dtype # torch.float64 + +Devices +------- + +- CUDA inputs stay on CUDA through ``fit``/``predict``/``transform`` — no + host round-trips, no device-mismatch errors. +- CPU remains the tested path: the suite runs on CPU, and CUDA coverage is + a ``requires CUDA``-gated round-trip test (see + ``tests/common/test_dtype_device.py``). + +.. code-block:: python + + X = torch.randn(40, 3, device="cuda") + StandardScaler().fit_transform(X).device.type # 'cuda' + +Caveat: on Apple Silicon, some ``torch.linalg`` ops used internally +(``lstsq``, ``svd``, ``cholesky``) have limited MPS support — CPU is the +supported backend there. diff --git a/tests/common/test_dtype_device.py b/tests/common/test_dtype_device.py new file mode 100644 index 0000000..d850a12 --- /dev/null +++ b/tests/common/test_dtype_device.py @@ -0,0 +1,132 @@ +"""Tests for dtype preservation and device propagation.""" + +from __future__ import annotations + +import pytest +import torch + +from torml.cluster import KMeans +from torml.decomposition import PCA +from torml.linear_model import LinearRegression +from torml.neighbors import KNeighborsClassifier +from torml.preprocessing import StandardScaler +from torml.tree import DecisionTreeClassifier +from torml.utils import check_array + +cuda_only = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" +) + + +class TestCheckArrayDtype: + """Tests for check_array dtype=None preservation.""" + + def test_preserves_float64(self): + """Test float64 tensors keep their dtype.""" + out = check_array(torch.randn(4, 2, dtype=torch.float64)) + assert out.dtype == torch.float64 + + def test_preserves_float32(self): + """Test float32 tensors keep their dtype.""" + out = check_array(torch.randn(4, 2, dtype=torch.float32)) + assert out.dtype == torch.float32 + + def test_promotes_int(self): + """Test integer tensors promote to float32.""" + out = check_array(torch.ones(4, 2, dtype=torch.int64)) + assert out.dtype == torch.float32 + + def test_list_defaults_float32(self): + """Test lists default to float32.""" + out = check_array([[1.0, 2.0], [3.0, 4.0]]) + assert out.dtype == torch.float32 + + def test_explicit_dtype_still_forces(self): + """Test explicit dtype overrides preservation.""" + out = check_array(torch.randn(4, 2, dtype=torch.float64), dtype=torch.float32) + assert out.dtype == torch.float32 + + +class TestDtypePreservation: + """Tests that estimators compute in the input dtype.""" + + def test_scaler_float64(self): + """Test StandardScaler learns float64 statistics.""" + X = torch.randn(20, 3, dtype=torch.float64) + sc = StandardScaler().fit(X) + assert sc.mean_.dtype == torch.float64 + assert sc.scale_.dtype == torch.float64 + assert sc.transform(X).dtype == torch.float64 + assert sc.inverse_transform(sc.transform(X)).dtype == torch.float64 + + def test_linear_regression_float64(self): + """Test LinearRegression learns float64 coefficients.""" + torch.manual_seed(0) + X = torch.randn(30, 2, dtype=torch.float64) + y = X[:, 0] * 2 - X[:, 1] + model = LinearRegression().fit(X, y) + assert model.coef_.dtype == torch.float64 + assert model.predict(X).dtype == torch.float64 + + def test_numeric_agreement(self): + """Test float32 and float64 runs agree closely.""" + torch.manual_seed(1) + X32 = torch.randn(30, 2) + y32 = X32[:, 0] * 2 - X32[:, 1] + X64, y64 = X32.double(), y32.double() + pred32 = LinearRegression().fit(X32, y32).predict(X32).double() + pred64 = LinearRegression().fit(X64, y64).predict(X64) + torch.testing.assert_close(pred32, pred64, rtol=1e-4, atol=1e-4) + + def test_int_features_accepted(self): + """Test integer features are promoted, not rejected.""" + X = torch.randint(0, 5, (20, 2)) + y = (X[:, 0] > 2).long() + pred = KNeighborsClassifier(3).fit(X, y).predict(X) + assert pred.shape == (20,) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) + def test_predict_accepts_both(self, dtype): + """Test classifiers predict on both float dtypes.""" + torch.manual_seed(2) + X = (torch.randn(20, 2)).to(dtype) + y = (X[:, 0] > 0).long() + for clf in (KNeighborsClassifier(3), DecisionTreeClassifier()): + assert clf.fit(X, y).predict(X).shape == (20,) + + def test_pca_float64(self): + """Test PCA keeps float64 through transform.""" + X = torch.randn(20, 4, dtype=torch.float64) + pca = PCA(n_components=2).fit(X) + assert pca.components_.dtype == torch.float64 + assert pca.transform(X).dtype == torch.float64 + + def test_kmeans_float64(self): + """Test KMeans centers match input dtype.""" + X = torch.randn(20, 2, dtype=torch.float64) + km = KMeans(n_clusters=2, random_state=0, n_init=2).fit(X) + assert km.cluster_centers_.dtype == torch.float64 + + +class TestDevice: + """Tests for device propagation (CPU always, CUDA when present).""" + + def test_cpu_round_trip(self): + """Test CPU tensors stay on CPU.""" + X = torch.randn(20, 2) + y = (X[:, 0] > 0).long() + assert StandardScaler().fit_transform(X).device.type == "cpu" + assert LinearRegression().fit(X, y.float()).predict(X).device.type == "cpu" + assert KMeans(n_clusters=2, random_state=0, n_init=2).fit(X).labels_.device.type == "cpu" + + @cuda_only + def test_cuda_round_trip(self): + """Test CUDA tensors stay on CUDA end to end.""" + torch.manual_seed(0) + X = torch.randn(40, 3, device="cuda") + y = (X[:, 0] > 0).long() + yr = X[:, 0] * 2 - X[:, 1] + assert StandardScaler().fit_transform(X).device.type == "cuda" + assert LinearRegression().fit(X, yr).predict(X).device.type == "cuda" + assert KNeighborsClassifier(3).fit(X, y).predict(X).device.type == "cuda" + assert KMeans(n_clusters=2, random_state=0, n_init=2).fit(X).labels_.device.type == "cuda" diff --git a/torml/base/_base.py b/torml/base/_base.py index 5ebb4b7..41606eb 100644 --- a/torml/base/_base.py +++ b/torml/base/_base.py @@ -186,7 +186,6 @@ def _check_X( return check_array( X, - dtype=torch.float32, ensure_2d=True, allow_nd=False, copy=True, @@ -204,7 +203,6 @@ def _check_y( return check_array( y, - dtype=torch.float32, ensure_2d=False, allow_nd=False, copy=True, diff --git a/torml/cluster/_dbscan.py b/torml/cluster/_dbscan.py index b845e73..d2a211a 100644 --- a/torml/cluster/_dbscan.py +++ b/torml/cluster/_dbscan.py @@ -78,7 +78,7 @@ def fit(self, X, y=None): ) if int(self.min_samples) < 1: raise ValueError(f"min_samples must be >= 1, got {self.min_samples}.") - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) n_samples = int(Xt.shape[0]) self.n_features_in_ = int(Xt.shape[1]) @@ -88,10 +88,12 @@ def fit(self, X, y=None): for i in range(n_samples) ] is_core = torch.tensor( - [len(nb) >= int(self.min_samples) for nb in neighborhoods], dtype=torch.bool + [len(nb) >= int(self.min_samples) for nb in neighborhoods], + dtype=torch.bool, + device=Xt.device, ) - labels = torch.full((n_samples,), -1, dtype=torch.long) - visited = torch.zeros(n_samples, dtype=torch.bool) + labels = torch.full((n_samples,), -1, dtype=torch.long, device=Xt.device) + visited = torch.zeros(n_samples, dtype=torch.bool, device=Xt.device) cluster_id = 0 for i in range(n_samples): if bool(visited[i]) or not bool(is_core[i]): @@ -119,5 +121,5 @@ def fit_predict(self, X, y=None): labels : torch.Tensor of shape (n_samples,) Cluster labels (-1 for noise). """ - check_array(X, ensure_2d=True, dtype=torch.float32) + check_array(X, ensure_2d=True) return self.fit(X, y).labels_ diff --git a/torml/cluster/_kmeans.py b/torml/cluster/_kmeans.py index df50166..59626f6 100644 --- a/torml/cluster/_kmeans.py +++ b/torml/cluster/_kmeans.py @@ -97,7 +97,7 @@ def fit(self, X, y=None): """ self._validate_hyperparams() generator = check_random_state(self.random_state) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) n_samples = int(Xt.shape[0]) if n_samples < int(self.n_clusters): raise ValueError( @@ -107,9 +107,9 @@ def fit(self, X, y=None): best_inertia: float | None = None for _ in range(int(self.n_init)): - perm = torch.randperm(n_samples, generator=generator) + perm = torch.randperm(n_samples, generator=generator, device=Xt.device) centers = Xt[perm[: int(self.n_clusters)]].clone() - assign = torch.zeros(n_samples, dtype=torch.long) + assign = torch.zeros(n_samples, dtype=torch.long, device=Xt.device) n_iter = 0 for n_iter in range(1, int(self.max_iter) + 1): dist = torch.cdist(Xt, centers, p=2) @@ -148,7 +148,7 @@ def predict(self, X): Closest-center indices. """ check_is_fitted(self, attributes=["cluster_centers_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but KMeans was fitted " diff --git a/torml/covariance/_empirical.py b/torml/covariance/_empirical.py index 46c1c3b..f275c05 100644 --- a/torml/covariance/_empirical.py +++ b/torml/covariance/_empirical.py @@ -45,7 +45,7 @@ def fit(self, X, y=None): self : EmpiricalCovariance Fitted estimator. """ - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) self.location_ = Xt.mean(dim=0) centered = Xt - self.location_ self.covariance_ = (centered.T @ centered) / int(Xt.shape[0]) @@ -69,7 +69,7 @@ def mahalanobis(self, X): Squared distances. """ check_is_fitted(self, attributes=["precision_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but the estimator was " @@ -94,10 +94,16 @@ def score(self, X, y=None): Mean log likelihood. """ check_is_fitted(self, attributes=["precision_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) d = int(Xt.shape[1]) - sign, log_det = torch.slogdet(self.covariance_ + 1e-12 * torch.eye(d)) + sign, log_det = torch.slogdet( + self.covariance_ + + 1e-12 + * torch.eye(d, dtype=self.covariance_.dtype, device=self.covariance_.device) + ) if int(sign.item()) <= 0: - log_det = torch.tensor(float("-inf")) + log_det = torch.tensor( + float("-inf"), dtype=log_det.dtype, device=log_det.device + ) ll = -0.5 * (d * math.log(2.0 * math.pi) + log_det + self.mahalanobis(Xt)) return float(ll.mean()) diff --git a/torml/cross_decomposition/_pls.py b/torml/cross_decomposition/_pls.py index b499743..a2bed4a 100644 --- a/torml/cross_decomposition/_pls.py +++ b/torml/cross_decomposition/_pls.py @@ -66,8 +66,9 @@ def fit(self, X, y): if int(self.n_components) < 1: raise ValueError(f"n_components must be >= 1, got {self.n_components}.") X, y = check_X_y(X, y) - X = X.to(dtype=torch.float32) - target = y.to(dtype=torch.float32).reshape(-1) + _dtype = X.dtype + X = X.to(dtype=_dtype) + target = y.to(dtype=_dtype).reshape(-1) n, d = int(X.shape[0]), int(X.shape[1]) k = min(int(self.n_components), d) self.n_features_in_ = d @@ -94,15 +95,15 @@ def fit(self, X, y): Xk = Xk - torch.outer(t, p) yk = yk - t * q if not weights: - self.x_weights_ = torch.zeros(d, 0, dtype=torch.float32) - self.x_scores_ = torch.zeros(n, 0, dtype=torch.float32) - self.coef_ = torch.zeros(d, dtype=torch.float32) + self.x_weights_ = torch.zeros(d, 0, dtype=X.dtype, device=X.device) + self.x_scores_ = torch.zeros(n, 0, dtype=X.dtype, device=X.device) + self.coef_ = torch.zeros(d, dtype=X.dtype, device=X.device) self.intercept_ = self._y_mean.reshape(1) self.n_components_ = 0 return self W = torch.stack(weights, dim=1) P = torch.stack([p for p, _ in loadings], dim=1) - qvec = torch.tensor([q for _, q in loadings], dtype=torch.float32) + qvec = torch.tensor([q for _, q in loadings], dtype=X.dtype, device=X.device) self.x_weights_ = W self.x_scores_ = torch.stack(scores, dim=1) self.coef_ = W @ torch.linalg.solve( # pylint: disable=not-callable @@ -126,7 +127,7 @@ def predict(self, X): Predicted values. """ check_is_fitted(self, attributes=["coef_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but PLSRegression was " @@ -148,5 +149,5 @@ def transform(self, X): Latent scores. """ check_is_fitted(self, attributes=["x_weights_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) return (Xt - self._x_mean) @ self.x_weights_ diff --git a/torml/decomposition/_pca.py b/torml/decomposition/_pca.py index 5b6103d..f9a3045 100644 --- a/torml/decomposition/_pca.py +++ b/torml/decomposition/_pca.py @@ -78,7 +78,7 @@ def fit(self, X, y=None): self : PCA Fitted transformer. """ - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) n_samples, n_features = int(Xt.shape[0]), int(Xt.shape[1]) if n_samples < 2: raise ValueError(f"Need at least 2 samples, got {n_samples}.") @@ -127,7 +127,7 @@ def transform(self, X): Projected data. """ check_is_fitted(self, attributes=["components_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but PCA was fitted " @@ -152,7 +152,7 @@ def inverse_transform(self, X): Approximate reconstruction. """ check_is_fitted(self, attributes=["components_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32, ensure_min_features=1) + Xt = check_array(X, ensure_2d=True, ensure_min_features=1) if int(Xt.shape[1]) != int(self.components_.shape[0]): raise ValueError( f"X has {int(Xt.shape[1])} components, but PCA has " diff --git a/torml/discriminant_analysis/_lda.py b/torml/discriminant_analysis/_lda.py index f0ac348..b588959 100644 --- a/torml/discriminant_analysis/_lda.py +++ b/torml/discriminant_analysis/_lda.py @@ -58,16 +58,16 @@ def fit(self, X, y): """ # pylint: disable=too-many-locals X, y = check_X_y(X, y) - X = X.to(dtype=torch.float32) flat = y.reshape(-1) self.classes_, inverse = torch.unique(flat, sorted=True, return_inverse=True) - inverse = inverse.to(dtype=torch.long) + self.classes_ = self.classes_.to(device=X.device) + inverse = inverse.to(dtype=torch.long, device=X.device) n, d = int(X.shape[0]), int(X.shape[1]) n_classes = int(self.classes_.shape[0]) if n_classes < 2: raise ValueError(f"Need at least 2 classes, got {n_classes}.") self.n_features_in_ = d - counts = torch.bincount(inverse, minlength=n_classes).to(dtype=torch.float32) + counts = torch.bincount(inverse, minlength=n_classes).to(dtype=X.dtype) self.priors_ = counts / counts.sum() self.means_ = torch.stack( [X[inverse == c].mean(dim=0) for c in range(n_classes)] @@ -80,7 +80,7 @@ def fit(self, X, y): V, S = Vt.T[:, :rank], singular[:rank] X_orth = (centered @ V) / S.clamp(min=1e-12) grand = X_orth.mean(dim=0) - between = torch.zeros(rank, rank, dtype=torch.float32) + between = torch.zeros(rank, rank, dtype=X.dtype, device=X.device) for c in range(n_classes): diff = X_orth[inverse == c].mean(dim=0) - grand between = between + counts[c] * torch.outer(diff, diff) @@ -109,7 +109,9 @@ def fit(self, X, y): self.scalings_ = scalings total = float(evals[:max_comp].sum()) self.explained_variance_ratio_ = ( - evals[:k] / total if total > 0 else torch.zeros(k) + evals[:k] / total + if total > 0 + else torch.zeros(k, dtype=X.dtype, device=X.device) ) self._scalings_full = (V / S.clamp(min=1e-12)) @ evecs[:, :max_comp] proj_means = self.means_ @ self._scalings_full @@ -121,7 +123,7 @@ def fit(self, X, y): mid @ direction ) self.coef_ = w - self.intercept_ = torch.tensor([b]) + self.intercept_ = torch.tensor([b], dtype=X.dtype, device=X.device) else: coefs, intercepts = [], [] for c in range(n_classes): @@ -132,13 +134,13 @@ def fit(self, X, y): coefs.append(w) intercepts.append(b) self.coef_ = torch.stack(coefs, dim=1) - self.intercept_ = torch.tensor(intercepts) + self.intercept_ = torch.tensor(intercepts, dtype=X.dtype, device=X.device) return self def _scores(self, X): """Unnormalized log posteriors for ``X``.""" check_is_fitted(self, attributes=["coef_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but LDA was fitted " @@ -201,7 +203,7 @@ def transform(self, X): Projected data. """ check_is_fitted(self, attributes=["scalings_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) return Xt @ self.scalings_ def _transform(self, X): diff --git a/torml/ensemble/_bagging.py b/torml/ensemble/_bagging.py index b810b99..1e5c862 100644 --- a/torml/ensemble/_bagging.py +++ b/torml/ensemble/_bagging.py @@ -51,9 +51,13 @@ def _fit_members(self, X, y, bootstrap: bool): members = [] for _ in range(int(self.n_estimators)): if bootstrap: - idx = torch.randint(n_samples, (k,), generator=generator) + idx = torch.randint( + n_samples, (k,), generator=generator, device=Xt.device + ) else: - idx = torch.randperm(n_samples, generator=generator)[:k] + idx = torch.randperm(n_samples, generator=generator, device=Xt.device)[ + :k + ] members.append((Xt, idx)) return members, generator @@ -168,7 +172,14 @@ def predict_proba(self, X): if isinstance(est.classes_, torch.Tensor) else list(est.classes_) ) - aligned = proba[:, torch.tensor([pos[c] for c in other], dtype=torch.long)] + aligned = proba[ + :, + torch.tensor( + [pos[c] for c in other], + dtype=torch.long, + device=proba.device if isinstance(proba, torch.Tensor) else None, + ), + ] total = aligned if total is None else total + aligned return total / len(self.estimators_) @@ -193,6 +204,11 @@ def predict(self, X): ) pos = {c: i for i, c in enumerate(own)} cols = [] + device = ( + self.classes_.device + if isinstance(self.classes_, torch.Tensor) + else (X.device if isinstance(X, torch.Tensor) else None) + ) for est in self.estimators_: pred = est.predict(X) vals = pred.tolist() if isinstance(pred, torch.Tensor) else list(pred) @@ -200,6 +216,7 @@ def predict(self, X): torch.tensor( [pos[v.item() if isinstance(v, torch.Tensor) else v] for v in vals], dtype=torch.long, + device=device, ) ) idx = torch.mode(torch.stack(cols), dim=0).values @@ -273,10 +290,11 @@ def fit(self, X, y): base = self._base() members, _ = self._fit_members(X, y, bool(self.bootstrap)) Xt = torch.as_tensor(X) if not isinstance(X, torch.Tensor) else X + _dtype = Xt.dtype if Xt.is_floating_point() else torch.float32 yt = ( - torch.as_tensor(y, dtype=torch.float32) + torch.as_tensor(y, dtype=_dtype) if not isinstance(y, torch.Tensor) - else y.to(dtype=torch.float32) + else y.to(dtype=_dtype) ) flat = yt.reshape(-1) self.estimators_ = [clone(base).fit(Xt[idx], flat[idx]) for _, idx in members] @@ -298,9 +316,6 @@ def predict(self, X): """ check_is_fitted(self, attributes=["estimators_"]) stacked = torch.stack( - [ - torch.as_tensor(est.predict(X), dtype=torch.float32) - for est in self.estimators_ - ] + [torch.as_tensor(est.predict(X)) for est in self.estimators_] ) return stacked.mean(dim=0) diff --git a/torml/ensemble/_forest.py b/torml/ensemble/_forest.py index 70b9e48..7804bfa 100644 --- a/torml/ensemble/_forest.py +++ b/torml/ensemble/_forest.py @@ -76,7 +76,9 @@ def fit(self, X, y): self.classes_ = torch.as_tensor(uniq) if numeric else uniq self.estimators_ = [] for _ in range(int(self.n_estimators)): - idx = torch.randint(n_samples, (n_samples,), generator=generator) + idx = torch.randint( + n_samples, (n_samples,), generator=generator, device=Xt.device + ) tree = DecisionTreeClassifier(max_depth=self.max_depth) self.estimators_.append(tree.fit(Xt[idx], flat[idx])) self.n_features_in_ = int(Xt.shape[1]) if Xt.ndim == 2 else 0 @@ -103,6 +105,11 @@ def predict(self, X): ) pos = {c: i for i, c in enumerate(own)} cols = [] + device = ( + self.classes_.device + if isinstance(self.classes_, torch.Tensor) + else (X.device if isinstance(X, torch.Tensor) else None) + ) for est in self.estimators_: pred = est.predict(X) vals = pred.tolist() if isinstance(pred, torch.Tensor) else list(pred) @@ -110,6 +117,7 @@ def predict(self, X): torch.tensor( [pos[v.item() if isinstance(v, torch.Tensor) else v] for v in vals], dtype=torch.long, + device=device, ) ) idx = torch.mode(torch.stack(cols), dim=0).values @@ -169,16 +177,19 @@ def fit(self, X, y): raise ValueError(f"n_estimators must be >= 1, got {self.n_estimators}.") generator = check_random_state(self.random_state) Xt = torch.as_tensor(X) if not isinstance(X, torch.Tensor) else X + _dtype = Xt.dtype if Xt.is_floating_point() else torch.float32 yt = ( - torch.as_tensor(y, dtype=torch.float32) + torch.as_tensor(y, dtype=_dtype) if not isinstance(y, torch.Tensor) - else y.to(dtype=torch.float32) + else y.to(dtype=_dtype) ) flat = yt.reshape(-1) n_samples = int(Xt.shape[0]) self.estimators_ = [] for _ in range(int(self.n_estimators)): - idx = torch.randint(n_samples, (n_samples,), generator=generator) + idx = torch.randint( + n_samples, (n_samples,), generator=generator, device=Xt.device + ) tree = DecisionTreeRegressor(max_depth=self.max_depth) self.estimators_.append(tree.fit(Xt[idx], flat[idx])) self.n_features_in_ = int(Xt.shape[1]) if Xt.ndim == 2 else 0 @@ -199,10 +210,7 @@ def predict(self, X): """ check_is_fitted(self, attributes=["estimators_"]) stacked = torch.stack( - [ - torch.as_tensor(est.predict(X), dtype=torch.float32) - for est in self.estimators_ - ] + [torch.as_tensor(est.predict(X)) for est in self.estimators_] ) return stacked.mean(dim=0) diff --git a/torml/ensemble/_voting.py b/torml/ensemble/_voting.py index 5bbb660..91f30a3 100644 --- a/torml/ensemble/_voting.py +++ b/torml/ensemble/_voting.py @@ -161,7 +161,8 @@ def _align_proba(self, proba, member_classes): pos = {c: i for i, c in enumerate(other)} for c in own: index.append(pos[c]) - return proba[:, torch.tensor(index, dtype=torch.long)] + device = proba.device if isinstance(proba, torch.Tensor) else None + return proba[:, torch.tensor(index, dtype=torch.long, device=device)] def predict_proba(self, X): """Average member probabilities. @@ -213,6 +214,11 @@ def predict(self, X): ) pos = {c: i for i, c in enumerate(own)} cols = [] + device = ( + self.classes_.device + if isinstance(self.classes_, torch.Tensor) + else (X.device if isinstance(X, torch.Tensor) else None) + ) for _, est in self.estimators_: pred = est.predict(X) vals = pred.tolist() if isinstance(pred, torch.Tensor) else list(pred) @@ -223,6 +229,7 @@ def predict(self, X): for v in vals ], dtype=torch.long, + device=device, ) ) idx = torch.mode(torch.stack(cols), dim=0).values @@ -324,9 +331,6 @@ def predict(self, X): check_is_fitted(self, attributes=["estimators_"]) stacked = torch.stack( - [ - torch.as_tensor(est.predict(X), dtype=torch.float32) - for _, est in self.estimators_ - ] + [torch.as_tensor(est.predict(X)) for _, est in self.estimators_] ) return stacked.mean(dim=0) diff --git a/torml/feature_extraction/_dict_vectorizer.py b/torml/feature_extraction/_dict_vectorizer.py index 510c730..cb8b8d9 100644 --- a/torml/feature_extraction/_dict_vectorizer.py +++ b/torml/feature_extraction/_dict_vectorizer.py @@ -111,7 +111,10 @@ def inverse_transform(self, X): Nonzero entries per row. """ check_is_fitted(self, attributes=["vocabulary_"]) - Xt = torch.as_tensor(X, dtype=torch.float32) + if isinstance(X, torch.Tensor) and X.is_floating_point(): + Xt = torch.as_tensor(X) + else: + Xt = torch.as_tensor(X, dtype=torch.float32) if int(Xt.shape[1]) != len(self.feature_names_): raise ValueError( f"X has {int(Xt.shape[1])} columns, but the vocabulary has " diff --git a/torml/feature_selection/_univariate.py b/torml/feature_selection/_univariate.py index 34014ec..fd71238 100644 --- a/torml/feature_selection/_univariate.py +++ b/torml/feature_selection/_univariate.py @@ -73,8 +73,8 @@ def _f_pvalue(score: float, df_between: int, df_within: int) -> float: def _group_scatter(Xt, flat, classes, overall): """Between/within sum of squares per feature for class groups.""" d = int(Xt.shape[1]) - ss_between = torch.zeros(d) - ss_within = torch.zeros(d) + ss_between = torch.zeros(d, dtype=Xt.dtype, device=Xt.device) + ss_within = torch.zeros(d, dtype=Xt.dtype, device=Xt.device) for c in classes.tolist(): group = Xt[flat == c] n_c = int(group.shape[0]) @@ -102,7 +102,7 @@ def f_classif(X, y): Survival probabilities under the F distribution. """ Xt, yt = check_X_y(X, y) - Xt = Xt.to(dtype=torch.float32) + Xt = Xt.to(dtype=Xt.dtype) flat = yt.reshape(-1) classes = torch.unique(flat, sorted=True) n = int(Xt.shape[0]) @@ -124,7 +124,8 @@ def f_classif(X, y): # P(X > x) for F(df_between, df_within) via the incomplete beta. pvalues = torch.tensor( [_f_pvalue(float(s), df_between, df_within) for s in scores.tolist()], - dtype=torch.float32, + dtype=Xt.dtype, + device=Xt.device, ) return scores, pvalues @@ -188,8 +189,8 @@ def fit(self, X, y): else: k = int(self.k) scores, pvalues = self.score_func(Xt, yt) - scores = torch.as_tensor(scores, dtype=torch.float32).reshape(-1) - pvalues = torch.as_tensor(pvalues, dtype=torch.float32).reshape(-1) + scores = torch.as_tensor(scores, dtype=Xt.dtype, device=Xt.device).reshape(-1) + pvalues = torch.as_tensor(pvalues, dtype=Xt.dtype, device=Xt.device).reshape(-1) if int(scores.shape[0]) != n_features: raise ValueError("score_func returned wrong number of scores.") self.scores_ = scores @@ -197,7 +198,7 @@ def fit(self, X, y): order = torch.argsort( torch.nan_to_num(scores, nan=float("-inf")), descending=True, stable=True ) - mask = torch.zeros(n_features, dtype=torch.bool) + mask = torch.zeros(n_features, dtype=torch.bool, device=Xt.device) mask[order[:k]] = True self._mask = mask return self @@ -233,7 +234,7 @@ def transform(self, X): Selected columns. """ check_is_fitted(self, attributes=["_mask"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but SelectKBest was " diff --git a/torml/gaussian_process/_gpr.py b/torml/gaussian_process/_gpr.py index b9829aa..b464b0a 100644 --- a/torml/gaussian_process/_gpr.py +++ b/torml/gaussian_process/_gpr.py @@ -71,13 +71,12 @@ def fit(self, X, y): if float(self.alpha) < 0: raise ValueError(f"alpha must be >= 0, got {self.alpha}.") X, y = check_X_y(X, y) - X = X.to(dtype=torch.float32) self.X_train_ = X - self.y_train_ = y.to(dtype=torch.float32).reshape(-1) + self.y_train_ = y.to(dtype=X.dtype, device=X.device).reshape(-1) self.n_features_in_ = int(X.shape[1]) noise = float(self.alpha) + 1e-10 base = self._rbf(X, X) - eye = torch.eye(int(X.shape[0])) + eye = torch.eye(int(X.shape[0]), dtype=X.dtype, device=X.device) jitter = noise for _ in range(5): try: @@ -112,7 +111,7 @@ def predict(self, X, return_std=False): Posterior standard deviations. """ check_is_fitted(self, attributes=["L_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but the process was " diff --git a/torml/linear_model.py b/torml/linear_model.py index f45fe58..97dce28 100644 --- a/torml/linear_model.py +++ b/torml/linear_model.py @@ -88,8 +88,7 @@ def fit(self, X: torch.Tensor, y: torch.Tensor): def _fit(self, X: torch.Tensor, y: torch.Tensor): """Fit linear model with linear least squares.""" X, y = check_X_y(X, y) - X = X.to(dtype=torch.float32) - y = y.to(dtype=torch.float32).reshape(-1, 1) + y = y.to(dtype=X.dtype).reshape(-1, 1) self.n_features_in_ = X.shape[1] if self.fit_intercept: @@ -125,7 +124,7 @@ def predict(self, X: torch.Tensor): check_is_fitted(self, attributes=["coef_", "intercept_"]) from torml.utils._validation import check_array - X = check_array(X, ensure_2d=True, dtype=torch.float32) + X = check_array(X, ensure_2d=True) if X.shape[1] != self.n_features_in_: raise ValueError( f"X has {X.shape[1]} features, but {type(self).__name__} " @@ -262,7 +261,6 @@ def fit(self, X: torch.Tensor, y: torch.Tensor): if int(self.max_iter) < 1: raise ValueError(f"max_iter must be >= 1, got {self.max_iter}.") X, y = check_X_y(X, y) - X = X.to(dtype=torch.float32) flat = y.reshape(-1) classes, inverse = torch.unique(flat, sorted=True, return_inverse=True) if int(classes.shape[0]) != 2: @@ -271,12 +269,12 @@ def fit(self, X: torch.Tensor, y: torch.Tensor): f"got {int(classes.shape[0])} classes." ) self.classes_ = classes - target = inverse.to(dtype=torch.float32) + target = inverse.to(dtype=X.dtype) n_samples, n_features = int(X.shape[0]), int(X.shape[1]) self.n_features_in_ = n_features lam = 1.0 / (n_samples * float(self.C)) - w = torch.zeros(n_features, dtype=torch.float32) - b = torch.tensor(0.0) + w = torch.zeros(n_features, dtype=X.dtype, device=X.device) + b = torch.tensor(0.0, dtype=X.dtype, device=X.device) for _ in range(int(self.max_iter)): logits = X @ w + b prob = torch.sigmoid(logits) @@ -304,7 +302,7 @@ def decision_function(self, X: torch.Tensor) -> torch.Tensor: check_is_fitted(self, attributes=["coef_", "intercept_"]) from torml.utils._validation import check_array - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but LogisticRegression " diff --git a/torml/manifold/_mds.py b/torml/manifold/_mds.py index 44b47a3..bfb46eb 100644 --- a/torml/manifold/_mds.py +++ b/torml/manifold/_mds.py @@ -62,7 +62,7 @@ def fit(self, X, y=None): "dissimilarity must be 'euclidean' or 'precomputed', " f"got {self.dissimilarity!r}." ) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) n_samples = int(Xt.shape[0]) if int(self.n_components) > n_samples: raise ValueError( diff --git a/torml/mixture/_gmm.py b/torml/mixture/_gmm.py index b0bbb3f..1e89331 100644 --- a/torml/mixture/_gmm.py +++ b/torml/mixture/_gmm.py @@ -52,7 +52,7 @@ def __init__(self, n_components=1, max_iter=100, tol=1e-3, random_state=None): def _estimate_log_gaussian(self, X): """Log density of each sample under each component.""" n, d = int(X.shape[0]), int(X.shape[1]) - out = torch.empty(n, self._k, dtype=torch.float32) + out = torch.empty(n, self._k, dtype=X.dtype, device=X.device) for c in range(self._k): diff = X - self.means_[c] try: @@ -98,7 +98,7 @@ def fit(self, X, y=None): if int(self.max_iter) < 1: raise ValueError(f"max_iter must be >= 1, got {self.max_iter}.") generator = check_random_state(self.random_state) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) n, d = int(Xt.shape[0]), int(Xt.shape[1]) if n < int(self.n_components): raise ValueError( @@ -107,7 +107,12 @@ def fit(self, X, y=None): self._k = int(self.n_components) self.n_features_in_ = d - resp = torch.rand(n, self._k, generator=generator) + 1e-6 + resp = ( + torch.rand( + n, self._k, generator=generator, dtype=Xt.dtype, device=Xt.device + ) + + 1e-6 + ) resp = resp / resp.sum(dim=1, keepdim=True) prev_ll: float | None = None self.converged_ = False @@ -115,11 +120,13 @@ def fit(self, X, y=None): nk = resp.sum(dim=0) + 1e-10 self.weights_ = nk / n self.means_ = (resp.T @ Xt) / nk.unsqueeze(1) - covs = torch.empty(self._k, d, d, dtype=torch.float32) + covs = torch.empty(self._k, d, d, dtype=Xt.dtype, device=Xt.device) for c in range(self._k): diff = Xt - self.means_[c] covs[c] = (diff * resp[:, c].unsqueeze(1)).T @ diff / nk[c] - covs[c] = covs[c] + 1e-6 * torch.eye(d) + covs[c] = covs[c] + 1e-6 * torch.eye( + d, dtype=Xt.dtype, device=Xt.device + ) self.covariances_ = covs log_resp = self._estimate_log_gaussian(Xt) + torch.log( self.weights_ @@ -147,7 +154,7 @@ def predict_proba(self, X): Responsibilities. """ check_is_fitted(self, attributes=["weights_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but the mixture was " @@ -206,7 +213,7 @@ def score(self, X, y=None): Mean log likelihood. """ check_is_fitted(self, attributes=["weights_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) log_resp = self._estimate_log_gaussian(Xt) + torch.log(self.weights_).unsqueeze( 0 ) diff --git a/torml/model_selection/_split.py b/torml/model_selection/_split.py index 3b0c12a..d54d4d5 100644 --- a/torml/model_selection/_split.py +++ b/torml/model_selection/_split.py @@ -130,9 +130,11 @@ def _to_n(size, name: str) -> int | None: raise ValueError("Train and test sets must each have at least 1 sample.") if shuffle: - perm = torch.randperm(n_samples, generator=generator) + _device = arrays[0].device if isinstance(arrays[0], torch.Tensor) else None + perm = torch.randperm(n_samples, generator=generator, device=_device) else: - perm = torch.arange(n_samples) + _device = arrays[0].device if isinstance(arrays[0], torch.Tensor) else None + perm = torch.arange(n_samples, device=_device) train_idx = perm[:n_train] test_idx = perm[n_train : n_train + n_test] @@ -240,9 +242,11 @@ def split(self, X, y=None, groups=None): ) if self.shuffle: - indices = torch.randperm(n_samples, generator=generator) + _device = X.device if isinstance(X, torch.Tensor) else None + indices = torch.randperm(n_samples, generator=generator, device=_device) else: - indices = torch.arange(n_samples) + _device = X.device if isinstance(X, torch.Tensor) else None + indices = torch.arange(n_samples, device=_device) fold_sizes = [n_samples // self.n_splits] * self.n_splits for i in range(n_samples % self.n_splits): diff --git a/torml/multiclass/_ovr.py b/torml/multiclass/_ovr.py index af9a839..ec3792a 100644 --- a/torml/multiclass/_ovr.py +++ b/torml/multiclass/_ovr.py @@ -76,10 +76,11 @@ def fit(self, X, y): self.classes_ = torch.as_tensor(uniq) if numeric else uniq self.estimators_ = [] for c in uniq: + cond = torch.as_tensor([v == c for v in flat.tolist()], device=Xt.device) binary = torch.where( - torch.as_tensor([v == c for v in flat.tolist()]), - torch.tensor(1), - torch.tensor(0), + cond, + torch.tensor(1, device=Xt.device), + torch.tensor(0, device=Xt.device), ) self.estimators_.append(clone(self.estimator).fit(X, binary)) first = self.estimators_[0] @@ -107,11 +108,9 @@ def decision_function(self, X): cols = [] for est in self.estimators_: if hasattr(est, "decision_function"): - s = torch.as_tensor( - est.decision_function(X), dtype=torch.float32 - ).reshape(-1) + s = torch.as_tensor(est.decision_function(X)).reshape(-1) elif hasattr(est, "predict_proba"): - s = torch.as_tensor(est.predict_proba(X), dtype=torch.float32)[:, 1] + s = torch.as_tensor(est.predict_proba(X))[:, 1] else: raise AttributeError("Members need decision_function or predict_proba.") cols.append(s) diff --git a/torml/multivariate/_multioutput.py b/torml/multivariate/_multioutput.py index cd830d5..f6ddb5b 100644 --- a/torml/multivariate/_multioutput.py +++ b/torml/multivariate/_multioutput.py @@ -82,8 +82,8 @@ def fit(self, X, y): "estimator must be a BaseEstimator, " f"got {type(self.estimator).__name__}." ) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) - Yt = _validate_2d_y(y, "y").to(dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) + Yt = _validate_2d_y(y, "y").to(dtype=Xt.dtype) if int(Xt.shape[0]) != int(Yt.shape[0]): raise ValueError( f"X and y have inconsistent lengths: {int(Xt.shape[0])} " @@ -110,10 +110,9 @@ def predict(self, X): Predictions per output. """ check_is_fitted(self, attributes=["estimators_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) cols = [ - torch.as_tensor(est.predict(Xt), dtype=torch.float32).reshape(-1) - for est in self.estimators_ + torch.as_tensor(est.predict(Xt)).reshape(-1) for est in self.estimators_ ] stacked = torch.stack(cols, dim=1) return stacked.squeeze(1) if stacked.shape[1] == 1 else stacked @@ -186,7 +185,7 @@ def fit(self, X, y): "estimator must be a BaseEstimator, " f"got {type(self.estimator).__name__}." ) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) Yt = _validate_2d_y(y, "y") if int(Xt.shape[0]) != int(Yt.shape[0]): raise ValueError( @@ -217,7 +216,7 @@ def predict(self, X): Predicted labels per output. """ check_is_fitted(self, attributes=["estimators_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) cols = [] for est in self.estimators_: pred = est.predict(Xt) diff --git a/torml/naive_bayes/_gaussian.py b/torml/naive_bayes/_gaussian.py index 69c4f95..31f17c7 100644 --- a/torml/naive_bayes/_gaussian.py +++ b/torml/naive_bayes/_gaussian.py @@ -68,7 +68,7 @@ def fit(self, X, y): Fitted classifier. """ vs = self._validate_smoothing() - X = check_array(X, ensure_2d=True, dtype=torch.float32) + X = check_array(X, ensure_2d=True) if isinstance(y, torch.Tensor): flat = y.reshape(-1) n_y = int(flat.shape[0]) @@ -94,7 +94,8 @@ def fit(self, X, y): self.classes_, inverse = torch.unique( labels, sorted=True, return_inverse=True ) - inverse = inverse.to(dtype=torch.long) + self.classes_ = self.classes_.to(device=X.device) + inverse = inverse.to(dtype=torch.long, device=X.device) else: try: uniq = sorted(set(labels)) @@ -102,7 +103,9 @@ def fit(self, X, y): raise TypeError("Labels must be sortable.") from e self.classes_ = uniq mapping = {c: i for i, c in enumerate(uniq)} - inverse = torch.tensor([mapping[v] for v in labels], dtype=torch.long) + inverse = torch.tensor( + [mapping[v] for v in labels], dtype=torch.long, device=X.device + ) n_classes = ( len(self.classes_) if isinstance(self.classes_, list) @@ -113,11 +116,11 @@ def fit(self, X, y): self.n_samples_in_ = int(X.shape[0]) self.class_count_ = torch.bincount(inverse, minlength=n_classes).to( - dtype=torch.float32 + dtype=X.dtype ) self.class_prior_ = self.class_count_ / self.class_count_.sum() - self.theta_ = torch.empty(n_classes, n_features, dtype=torch.float32) - self.var_ = torch.empty(n_classes, n_features, dtype=torch.float32) + self.theta_ = torch.empty(n_classes, n_features, dtype=X.dtype, device=X.device) + self.var_ = torch.empty(n_classes, n_features, dtype=X.dtype, device=X.device) for c in range(n_classes): X_c = X[inverse == c] self.theta_[c] = X_c.mean(dim=0) @@ -130,7 +133,7 @@ def fit(self, X, y): def _joint_log_likelihood(self, X: torch.Tensor) -> torch.Tensor: """Compute unnormalized log posteriors for ``X``.""" check_is_fitted(self, attributes=["theta_", "var_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but GaussianNB was " diff --git a/torml/neighbors/_classification.py b/torml/neighbors/_classification.py index 9595a40..1003227 100644 --- a/torml/neighbors/_classification.py +++ b/torml/neighbors/_classification.py @@ -67,7 +67,6 @@ def fit(self, X, y): p = _validate_p(self.p) self._p = p X, y = check_X_y(X, y) - X = X.to(dtype=torch.float32) if int(X.shape[0]) < int(self.n_neighbors): raise ValueError( f"n_neighbors ({self.n_neighbors}) must be <= n_samples " @@ -101,7 +100,7 @@ def kneighbors(self, X, n_neighbors=None, return_distance=True): Indices into the training set. """ check_is_fitted(self, attributes=["X_", "y_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but " @@ -137,29 +136,33 @@ def predict_proba(self, X): n_queries = int(ind.shape[0]) n_classes = int(self.classes_.shape[0]) if self.weights == "uniform": - counts = torch.zeros(n_queries, n_classes, dtype=torch.float32) + counts = torch.zeros( + n_queries, n_classes, dtype=dist.dtype, device=dist.device + ) counts.scatter_add_( 1, neighbor_labels, - torch.ones_like(neighbor_labels, dtype=torch.float32), + torch.ones_like(neighbor_labels, dtype=dist.dtype), ) return counts / float(self.n_neighbors) weights = distance_weights(dist) exact = dist == 0 has_exact = exact.any(dim=1) - proba = torch.zeros(n_queries, n_classes, dtype=torch.float32) + proba = torch.zeros(n_queries, n_classes, dtype=dist.dtype, device=dist.device) if bool(has_exact.any()): for i in torch.where(has_exact)[0].tolist(): exact_labels = neighbor_labels[i][exact[i]] counts = torch.bincount(exact_labels, minlength=n_classes).to( - dtype=torch.float32 + dtype=dist.dtype ) proba[i] = counts / counts.sum() rest = (~has_exact).nonzero().reshape(-1).tolist() if rest: - r_idx = torch.tensor(rest, dtype=torch.long) + r_idx = torch.tensor(rest, dtype=torch.long, device=dist.device) w = weights[r_idx] - weighted = torch.zeros(len(rest), n_classes, dtype=torch.float32) + weighted = torch.zeros( + len(rest), n_classes, dtype=dist.dtype, device=dist.device + ) weighted.scatter_add_(1, neighbor_labels[r_idx], w) proba[r_idx] = weighted / w.sum(dim=1, keepdim=True) return proba diff --git a/torml/neighbors/_regression.py b/torml/neighbors/_regression.py index 983c5c3..2aaa197 100644 --- a/torml/neighbors/_regression.py +++ b/torml/neighbors/_regression.py @@ -65,8 +65,7 @@ def fit(self, X, y): p = _validate_p(self.p) self._p = p X, y = check_X_y(X, y) - X = X.to(dtype=torch.float32) - y = y.to(dtype=torch.float32).reshape(-1) + y = y.to(dtype=X.dtype).reshape(-1) if int(X.shape[0]) < int(self.n_neighbors): raise ValueError( f"n_neighbors ({self.n_neighbors}) must be <= n_samples " @@ -98,7 +97,7 @@ def kneighbors(self, X, n_neighbors=None, return_distance=True): Indices into the training set. """ check_is_fitted(self, attributes=["X_", "y_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but " @@ -135,13 +134,15 @@ def predict(self, X): return neighbor_vals.mean(dim=1) exact = dist == 0 has_exact = exact.any(dim=1) - out = torch.empty(int(ind.shape[0]), dtype=torch.float32) + out = torch.empty( + int(ind.shape[0]), dtype=neighbor_vals.dtype, device=neighbor_vals.device + ) if bool(has_exact.any()): for i in torch.where(has_exact)[0].tolist(): out[i] = neighbor_vals[i][exact[i]].mean() rest = (~has_exact).nonzero().reshape(-1).tolist() if rest: - r_idx = torch.tensor(rest, dtype=torch.long) + r_idx = torch.tensor(rest, dtype=torch.long, device=neighbor_vals.device) w = distance_weights(dist[r_idx]) out[r_idx] = (neighbor_vals[r_idx] * w).sum(dim=1) / w.sum(dim=1) return out diff --git a/torml/pipelines/_union.py b/torml/pipelines/_union.py index 55f8acb..1fe71a3 100644 --- a/torml/pipelines/_union.py +++ b/torml/pipelines/_union.py @@ -177,7 +177,7 @@ def fit(self, X, y=None): Fitted transformer. """ cleaned = self._validate() - arr = check_array(X, ensure_2d=True, dtype=torch.float32) + arr = check_array(X, ensure_2d=True) n_features = int(arr.shape[1]) self.n_features_in_ = n_features fitted = [] @@ -206,7 +206,7 @@ def transform(self, X): Concatenated outputs. """ check_is_fitted(self, attributes=["transformers_"]) - arr = check_array(X, ensure_2d=True, dtype=torch.float32) + arr = check_array(X, ensure_2d=True) if int(arr.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(arr.shape[1])} features, but ColumnTransformer " diff --git a/torml/preprocessing/_data.py b/torml/preprocessing/_data.py index b9a84e2..aa7b9f8 100644 --- a/torml/preprocessing/_data.py +++ b/torml/preprocessing/_data.py @@ -62,7 +62,7 @@ def fit(self, X, y=None): raise TypeError( f"with_std must be a bool, got {type(self.with_std).__name__}." ) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) self.n_features_in_ = int(Xt.shape[1]) self.n_samples_seen_ = int(Xt.shape[0]) if self.with_mean: @@ -99,7 +99,7 @@ def transform(self, X): Transformed data. """ check_is_fitted(self, attributes=["mean_", "scale_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but StandardScaler was " @@ -129,7 +129,7 @@ def inverse_transform(self, X): Data in the original space. """ check_is_fitted(self, attributes=["mean_", "scale_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but StandardScaler was " @@ -205,7 +205,7 @@ def fit(self, X, y=None): Fitted scaler. """ lo, hi = self._validate_range() - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) self.n_features_in_ = int(Xt.shape[1]) self.n_samples_seen_ = int(Xt.shape[0]) self.data_min_ = Xt.min(dim=0).values @@ -233,7 +233,7 @@ def transform(self, X): Transformed data. """ check_is_fitted(self, attributes=["scale_", "min_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but MinMaxScaler was " @@ -258,7 +258,7 @@ def inverse_transform(self, X): Data in the original space. """ check_is_fitted(self, attributes=["scale_", "min_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but MinMaxScaler was " diff --git a/torml/preprocessing/_encoders.py b/torml/preprocessing/_encoders.py index 250a516..97e3bbf 100644 --- a/torml/preprocessing/_encoders.py +++ b/torml/preprocessing/_encoders.py @@ -166,7 +166,13 @@ def transform(self, X): f"fitted with {int(self.n_features_in_)} features." ) n_samples = len(rows) - out = torch.zeros(n_samples, int(self.n_features_out_), dtype=self._out_dtype) + input_device = X.device if isinstance(X, torch.Tensor) else None + out = torch.zeros( + n_samples, + int(self.n_features_out_), + dtype=self._out_dtype, + device=input_device, + ) offset = 0 for j, width in enumerate(self._widths): mapping = self._cat_index[j] @@ -198,7 +204,7 @@ def inverse_transform(self, X): Decoded categorical rows. """ check_is_fitted(self, attributes=["categories_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_out_): raise ValueError( f"X has {int(Xt.shape[1])} columns, but OneHotEncoder output " diff --git a/torml/random_projection/_gaussian.py b/torml/random_projection/_gaussian.py index 4bbbce7..fbe58c4 100644 --- a/torml/random_projection/_gaussian.py +++ b/torml/random_projection/_gaussian.py @@ -82,7 +82,7 @@ def fit(self, X, y=None): Fitted projector. """ generator = check_random_state(self.random_state) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) n_samples, n_features = int(Xt.shape[0]), int(Xt.shape[1]) self.n_features_in_ = n_features nc = self.n_components @@ -100,7 +100,7 @@ def fit(self, X, y=None): k = int(nc) self.n_components_ = k self.components_ = torch.randn( - n_features, k, generator=generator, dtype=torch.float32 + n_features, k, generator=generator, dtype=Xt.dtype, device=Xt.device ) / math.sqrt(k) return self @@ -118,7 +118,7 @@ def transform(self, X): Projected data. """ check_is_fitted(self, attributes=["components_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but the projector was " @@ -143,7 +143,7 @@ def inverse_transform(self, X): Approximate reconstruction. """ check_is_fitted(self, attributes=["components_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_components_): raise ValueError( f"X has {int(Xt.shape[1])} components, but the projector has " diff --git a/torml/semi_supervised/_label_propagation.py b/torml/semi_supervised/_label_propagation.py index ae8d088..2c4e4c8 100644 --- a/torml/semi_supervised/_label_propagation.py +++ b/torml/semi_supervised/_label_propagation.py @@ -66,7 +66,7 @@ def fit(self, X, y): ) if int(self.n_neighbors) < 1: raise ValueError(f"n_neighbors must be >= 1, got {self.n_neighbors}.") - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) flat = ( torch.as_tensor(y).reshape(-1) if not isinstance(y, torch.Tensor) @@ -94,14 +94,14 @@ def fit(self, X, y): dist.fill_diagonal_(float("inf")) k = min(int(self.n_neighbors), n - 1) knn = torch.topk(dist, k=k, dim=1, largest=False).indices - weight = torch.zeros(n, n, dtype=torch.float32) + weight = torch.zeros(n, n, dtype=Xt.dtype, device=Xt.device) sigma = float(torch.median(dist[dist < float("inf")])) + 1e-12 for i in range(n): weight[i, knn[i]] = torch.exp(-dist[i, knn[i]] ** 2 / (2 * sigma**2)) weight = torch.maximum(weight, weight.T) weight = weight / weight.sum(dim=1, keepdim=True).clamp(min=1e-12) - Y = torch.zeros(n, n_classes, dtype=torch.float32) + Y = torch.zeros(n, n_classes, dtype=Xt.dtype, device=Xt.device) for i in torch.where(labeled_mask)[0].tolist(): v = flat[i].item() Y[i, pos[v]] = 1.0 @@ -134,7 +134,7 @@ def predict(self, X): Propagated labels. """ check_is_fitted(self, attributes=["transduction_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[0]) != int(self.transduction_.shape[0]): raise ValueError( "LabelPropagation is transductive: predict on the fitted X." @@ -155,7 +155,7 @@ def predict_proba(self, X): Soft assignments. """ check_is_fitted(self, attributes=["label_distributions_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[0]) != int(self.label_distributions_.shape[0]): raise ValueError( "LabelPropagation is transductive: predict on the fitted X." diff --git a/torml/svm/_linear.py b/torml/svm/_linear.py index 8eb5fc5..fff801e 100644 --- a/torml/svm/_linear.py +++ b/torml/svm/_linear.py @@ -89,7 +89,6 @@ def fit(self, X, y): max_iter = _validate_rate("max_iter", self.max_iter) generator = check_random_state(self.random_state) X, y = check_X_y(X, y) - X = X.to(dtype=torch.float32) flat = y.reshape(-1) classes, inverse = torch.unique(flat, sorted=True, return_inverse=True) if int(classes.shape[0]) != 2: @@ -98,16 +97,24 @@ def fit(self, X, y): f"got {int(classes.shape[0])} classes." ) self.classes_ = classes - signed = torch.where(inverse == 1, 1.0, -1.0) + signed = torch.where( + inverse == 1, + torch.tensor(1.0, dtype=X.dtype, device=X.device), + torch.tensor(-1.0, dtype=X.dtype, device=X.device), + ) n_samples, n_features = int(X.shape[0]), int(X.shape[1]) self.n_features_in_ = n_features lam = 1.0 / (n_samples * alpha) - w = torch.zeros(n_features, dtype=torch.float32) - b = torch.tensor(0.0) + w = torch.zeros(n_features, dtype=X.dtype, device=X.device) + b = torch.tensor(0.0, dtype=X.dtype, device=X.device) prev_obj: float | None = None n_iter = 0 for t in range(1, max_iter + 1): - i = int(torch.randint(n_samples, (1,), generator=generator).item()) + i = int( + torch.randint( + n_samples, (1,), generator=generator, device=X.device + ).item() + ) eta = 1.0 / (lam * t) margin = signed[i] * (torch.dot(w, X[i]) + b) w = (1.0 - eta * lam) * w @@ -142,7 +149,7 @@ def decision_function(self, X): Positive values predict ``classes_[1]``. """ check_is_fitted(self, attributes=["coef_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but LinearSVC was fitted " @@ -226,15 +233,18 @@ def fit(self, X, y): raise ValueError(f"epsilon must be >= 0, got {self.epsilon}.") generator = check_random_state(self.random_state) X, y = check_X_y(X, y) - X = X.to(dtype=torch.float32) - target = y.to(dtype=torch.float32).reshape(-1) + target = y.to(dtype=X.dtype, device=X.device).reshape(-1) n_samples, n_features = int(X.shape[0]), int(X.shape[1]) self.n_features_in_ = n_features lam = 1.0 / (n_samples * alpha) - w = torch.zeros(n_features, dtype=torch.float32) - b = torch.tensor(0.0) + w = torch.zeros(n_features, dtype=X.dtype, device=X.device) + b = torch.tensor(0.0, dtype=X.dtype, device=X.device) for t in range(1, max_iter + 1): - i = int(torch.randint(n_samples, (1,), generator=generator).item()) + i = int( + torch.randint( + n_samples, (1,), generator=generator, device=X.device + ).item() + ) eta = 1.0 / (lam * t) residual = torch.dot(w, X[i]) + b - target[i] w = (1.0 - eta * lam) * w @@ -263,7 +273,7 @@ def predict(self, X): Predicted values. """ check_is_fitted(self, attributes=["coef_"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but LinearSVR was fitted " diff --git a/torml/tree/_tree.py b/torml/tree/_tree.py index 1a4dac5..d951045 100644 --- a/torml/tree/_tree.py +++ b/torml/tree/_tree.py @@ -27,7 +27,7 @@ def _best_split_classifier(X, y_idx, n_classes, criterion): # pylint: disable=too-many-locals n, n_features = int(X.shape[0]), int(X.shape[1]) impurity_fn = _gini if criterion == "gini" else _entropy - total_counts = torch.bincount(y_idx, minlength=n_classes).to(dtype=torch.float32) + total_counts = torch.bincount(y_idx, minlength=n_classes).to(dtype=X.dtype) parent = impurity_fn(total_counts) best = (0.0, -1, 0.0) for f in range(n_features): @@ -37,10 +37,10 @@ def _best_split_classifier(X, y_idx, n_classes, criterion): valid = sorted_x[:-1] != sorted_x[1:] if not bool(valid.any()): continue - one_hot = torch.zeros(n, n_classes, dtype=torch.float32) - one_hot[torch.arange(n), y_idx[order]] = 1.0 + one_hot = torch.zeros(n, n_classes, dtype=X.dtype, device=X.device) + one_hot[torch.arange(n, device=X.device), y_idx[order]] = 1.0 left_counts = torch.cumsum(one_hot, dim=0)[:-1] - left_n = torch.arange(1, n, dtype=torch.float32) + left_n = torch.arange(1, n, dtype=X.dtype, device=X.device) right_counts = total_counts - left_counts right_n = float(n) - left_n gain = ( @@ -70,7 +70,7 @@ def _best_split_regressor(X, y): valid = sorted_x[:-1] != sorted_x[1:] if not bool(valid.any()): continue - left_n = torch.arange(1, n, dtype=torch.float32) + left_n = torch.arange(1, n, dtype=X.dtype, device=X.device) right_n = float(n) - left_n left_s = torch.cumsum(sorted_y, dim=0)[:-1] left_s2 = torch.cumsum(sorted_y * sorted_y, dim=0)[:-1] @@ -177,7 +177,6 @@ def fit(self, X, y): from torml.utils._validation import check_X_y as _check_X_y X, y = _check_X_y(X, y) - X = X.to(dtype=torch.float32) flat = y.reshape(-1) labels = flat.tolist() try: @@ -192,7 +191,9 @@ def fit(self, X, y): else uniq ) index_of = {c: i for i, c in enumerate(uniq)} - y_idx = torch.tensor([index_of[v] for v in labels], dtype=torch.long) + y_idx = torch.tensor( + [index_of[v] for v in labels], dtype=torch.long, device=X.device + ) n_classes = len(uniq) self.n_features_in_ = int(X.shape[1]) self._left, self._right, self._feature = [], [], [] @@ -202,15 +203,17 @@ def fit(self, X, y): [], [], ) - self._importance = torch.zeros(self.n_features_in_, dtype=torch.float32) + self._importance = torch.zeros( + self.n_features_in_, dtype=X.dtype, device=X.device + ) max_depth = float("inf") if self.max_depth is None else int(self.max_depth) - stack = [(0, torch.arange(int(X.shape[0])), 0)] + stack = [(0, torch.arange(int(X.shape[0]), device=X.device), 0)] self._new_node(0.0, int(X.shape[0]), None) while stack: node, idx, depth = stack.pop() node_y = y_idx[idx] - counts = torch.bincount(node_y, minlength=n_classes).to(dtype=torch.float32) + counts = torch.bincount(node_y, minlength=n_classes).to(dtype=X.dtype) imp = _gini(counts) if self.criterion == "gini" else _entropy(counts) self._impurity[node] = float(imp.item()) self._values[node] = counts @@ -269,13 +272,18 @@ def predict_proba(self, X): Per-class probabilities. """ check_is_fitted(self, attributes=["_left"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but the tree was fitted " f"with {int(self.n_features_in_)} features." ) - out = torch.empty(int(Xt.shape[0]), len(self._values[0]), dtype=torch.float32) + out = torch.empty( + int(Xt.shape[0]), + len(self._values[0]), + dtype=Xt.dtype, + device=Xt.device, + ) for i in range(int(Xt.shape[0])): counts = self._values[self._predict_idx(Xt[i])] out[i] = counts / counts.sum().clamp(min=1) @@ -350,8 +358,7 @@ def fit(self, X, y): from torml.utils._validation import check_X_y as _check_X_y X, y = _check_X_y(X, y) - X = X.to(dtype=torch.float32) - target = y.to(dtype=torch.float32).reshape(-1) + target = y.to(dtype=X.dtype).reshape(-1) self.n_features_in_ = int(X.shape[1]) self._left, self._right, self._feature = [], [], [] self._threshold, self._impurity, self._n_node_samples, self._values = ( @@ -360,10 +367,12 @@ def fit(self, X, y): [], [], ) - self._importance = torch.zeros(self.n_features_in_, dtype=torch.float32) + self._importance = torch.zeros( + self.n_features_in_, dtype=X.dtype, device=X.device + ) max_depth = float("inf") if self.max_depth is None else int(self.max_depth) - stack = [(0, torch.arange(int(X.shape[0])), 0)] + stack = [(0, torch.arange(int(X.shape[0]), device=X.device), 0)] self._new_node(0.0, int(X.shape[0]), 0.0) while stack: node, idx, depth = stack.pop() @@ -410,13 +419,13 @@ def predict(self, X): Predicted values. """ check_is_fitted(self, attributes=["_left"]) - Xt = check_array(X, ensure_2d=True, dtype=torch.float32) + Xt = check_array(X, ensure_2d=True) if int(Xt.shape[1]) != int(self.n_features_in_): raise ValueError( f"X has {int(Xt.shape[1])} features, but the tree was fitted " f"with {int(self.n_features_in_)} features." ) - out = torch.empty(int(Xt.shape[0]), dtype=torch.float32) + out = torch.empty(int(Xt.shape[0]), dtype=Xt.dtype, device=Xt.device) for i in range(int(Xt.shape[0])): node = 0 row = Xt[i] diff --git a/torml/utils/_mask.py b/torml/utils/_mask.py index 81a6080..298bf43 100644 --- a/torml/utils/_mask.py +++ b/torml/utils/_mask.py @@ -28,7 +28,6 @@ def safe_mask(X: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: mask = torch.as_tensor(mask, device=X.device) mask = check_array( mask, - dtype=torch.float32, ensure_2d=False, allow_nd=False, copy=False, @@ -73,6 +72,6 @@ def indices_to_mask(indices: torch.Tensor, n_samples: int): idx = torch.as_tensor(indices, dtype=torch.long).reshape(-1) if idx.numel() > 0 and (bool((idx < 0).any()) or bool((idx >= n_samples).any())): raise ValueError("indices contains out-of-bounds entries.") - mask = torch.zeros((n_samples,), dtype=torch.uint8) + mask = torch.zeros((n_samples,), dtype=torch.uint8, device=idx.device) mask[idx] = 1 return mask diff --git a/torml/utils/_validation.py b/torml/utils/_validation.py index 751ce18..9bd1e8c 100644 --- a/torml/utils/_validation.py +++ b/torml/utils/_validation.py @@ -34,7 +34,7 @@ class NotFittedError(ValueError): def check_array( array, *, - dtype=torch.float32, + dtype=None, ensure_2d: bool = True, allow_nd: bool = False, copy: bool = False, @@ -50,8 +50,10 @@ def check_array( ---------- array : array-like of shape (n_samples, n_features) or (n_samples,) Data to check. Can be a torch.Tensor, numpy.ndarray, or list. - dtype : data-type, default=torch.float32 - Data type for the output tensor. + dtype : data-type or None, default=None + Data type for the output tensor. None preserves floating-point + tensor dtypes (integer tensors are promoted to ``torch.float32``); + non-tensor inputs default to ``torch.float32``. ensure_2d : bool, default=True If True, X will be converted to 2D array. If False, 1D arrays are allowed. allow_nd : bool, default=False @@ -81,8 +83,14 @@ def check_array( ValueError If ``array`` does not pass validation (wrong shape, contains NaN/Inf, etc.). """ - # Convert to tensor - dtype = _normalize_dtype(dtype) + # Resolve dtype: preserve floating tensor dtypes, default otherwise. + if dtype is None: + if isinstance(array, torch.Tensor) and array.is_floating_point(): + dtype = array.dtype + else: + dtype = torch.float32 + else: + dtype = _normalize_dtype(dtype) if isinstance(array, torch.Tensor): tensor = array.to(dtype=dtype) if array.dtype != dtype else array else: From e93527d508170cca7e70548bbda01988c24e95d7 Mon Sep 17 00:00:00 2001 From: mskumar Date: Fri, 25 Sep 2026 22:46:19 +0530 Subject: [PATCH 2/2] fix: black formatting for dtype/device tests --- tests/common/test_dtype_device.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/common/test_dtype_device.py b/tests/common/test_dtype_device.py index d850a12..3d27de4 100644 --- a/tests/common/test_dtype_device.py +++ b/tests/common/test_dtype_device.py @@ -13,9 +13,7 @@ from torml.tree import DecisionTreeClassifier from torml.utils import check_array -cuda_only = pytest.mark.skipif( - not torch.cuda.is_available(), reason="requires CUDA" -) +cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") class TestCheckArrayDtype: @@ -117,7 +115,10 @@ def test_cpu_round_trip(self): y = (X[:, 0] > 0).long() assert StandardScaler().fit_transform(X).device.type == "cpu" assert LinearRegression().fit(X, y.float()).predict(X).device.type == "cpu" - assert KMeans(n_clusters=2, random_state=0, n_init=2).fit(X).labels_.device.type == "cpu" + assert ( + KMeans(n_clusters=2, random_state=0, n_init=2).fit(X).labels_.device.type + == "cpu" + ) @cuda_only def test_cuda_round_trip(self): @@ -129,4 +130,7 @@ def test_cuda_round_trip(self): assert StandardScaler().fit_transform(X).device.type == "cuda" assert LinearRegression().fit(X, yr).predict(X).device.type == "cuda" assert KNeighborsClassifier(3).fit(X, y).predict(X).device.type == "cuda" - assert KMeans(n_clusters=2, random_state=0, n_init=2).fit(X).labels_.device.type == "cuda" + assert ( + KMeans(n_clusters=2, random_state=0, n_init=2).fit(X).labels_.device.type + == "cuda" + )