Preserve dtypes and propagate CUDA/MPS devices; 10k-sample comparison

doc/user_guide/dtypes_devices.rst

@@ -25,16 +25,19 @@ 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``).
+  host round-trips, no device-mismatch errors. Random generators are created
+  on the data device, and class tables live next to the data they index.
+- Apple Silicon (MPS) is supported the same way and is exercised by gated
+  tests wherever MPS exists (``tests/common/test_dtype_device.py``).
+- CPU remains the fully tested path: the suite runs on CPU, and CUDA/MPS
+  coverage is ``requires CUDA``/``requires MPS``-gated.
 
 .. code-block:: python
 
-   X = torch.randn(40, 3, device="cuda")
+   X = torch.randn(40, 3, device="cuda")   # or "mps"
    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.
+Caveat: on Apple Silicon, a few ``torch.linalg`` ops historically lagged on
+MPS — if you hit ``NotImplementedError`` from inside ``torch.linalg``, move
+that step to CPU. Majority votes use a gather-based implementation instead
+of ``torch.mode`` for exactly this reason.

tests/common/test_dtype_device.py

@@ -14,6 +14,9 @@
 from torml.utils import check_array
 
 cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+mps_only = pytest.mark.skipif(
+    not torch.backends.mps.is_available(), reason="requires MPS"
+)
 
 
 class TestCheckArrayDtype:
@@ -134,3 +137,50 @@ def test_cuda_round_trip(self):
             KMeans(n_clusters=2, random_state=0, n_init=2).fit(X).labels_.device.type
             == "cuda"
         )
+
+
+class TestMPS:
+    """Tests for Apple Silicon propagation (run where MPS exists)."""
+
+    @mps_only
+    def test_estimators_stay_on_mps(self):
+        """Test core estimators stay on MPS end to end."""
+        from torml.ensemble import RandomForestClassifier
+        from torml.mixture import GaussianMixture
+        from torml.svm import LinearSVC
+
+        torch.manual_seed(0)
+        X = torch.randn(40, 3, device="mps")
+        y = (X[:, 0] > 0).long()
+        yr = X[:, 0] * 2 - X[:, 1]
+        assert StandardScaler().fit_transform(X).device.type == "mps"
+        assert LinearRegression().fit(X, yr).predict(X).device.type == "mps"
+        assert KNeighborsClassifier(3).fit(X, y).predict(X).device.type == "mps"
+        assert DecisionTreeClassifier().fit(X, y).predict(X).device.type == "mps"
+        assert (
+            KMeans(n_clusters=2, random_state=0, n_init=2).fit(X).labels_.device.type
+            == "mps"
+        )
+        assert (
+            GaussianMixture(n_components=2, random_state=0)
+            .fit(X)
+            .predict(X)
+            .device.type
+            == "mps"
+        )
+        assert LinearSVC(random_state=0).fit(X, y).predict(X).device.type == "mps"
+        assert (
+            RandomForestClassifier(n_estimators=2, random_state=0)
+            .fit(X, y)
+            .predict(X)
+            .device.type
+            == "mps"
+        )
+
+    @mps_only
+    def test_mps_label_tables(self):
+        """Test class tables work with on-device indices."""
+        from torml.preprocessing import LabelEncoder
+
+        y = torch.tensor([0, 1, 0], device="mps")
+        assert LabelEncoder().fit_transform(y).device.type == "mps"

torml/cluster/_kmeans.py

@@ -96,7 +96,8 @@ def fit(self, X, y=None):
             Fitted clusterer.
         """
         self._validate_hyperparams()
-        generator = check_random_state(self.random_state)
+        device = X.device if isinstance(X, torch.Tensor) else None
+        generator = check_random_state(self.random_state, device)
         Xt = check_array(X, ensure_2d=True)
         n_samples = int(Xt.shape[0])
         if n_samples < int(self.n_clusters):

torml/cross_decomposition/_pls.py

@@ -67,8 +67,8 @@ def fit(self, X, y):
             raise ValueError(f"n_components must be >= 1, got {self.n_components}.")
         X, y = check_X_y(X, y)
         _dtype = X.dtype
-        X = X.to(dtype=_dtype)
-        target = y.to(dtype=_dtype).reshape(-1)
+        X = X.to(dtype=_dtype, device=X.device)
+        target = y.to(dtype=_dtype, device=X.device).reshape(-1)
         n, d = int(X.shape[0]), int(X.shape[1])
         k = min(int(self.n_components), d)
         self.n_features_in_ = d

torml/ensemble/_bagging.py

@@ -5,6 +5,7 @@
 import torch
 
 from torml.base import BaseEstimator, ClassifierMixin, RegressorMixin, clone
+from torml.ensemble._voting import _majority_vote
 from torml.utils._random import check_random_state
 from torml.utils._validation import check_is_fitted
 
@@ -44,7 +45,8 @@ def _fit_members(self, X, y, bootstrap: bool):
             )
         if int(self.n_estimators) < 1:
             raise ValueError(f"n_estimators must be >= 1, got {self.n_estimators}.")
-        generator = check_random_state(self.random_state)
+        device = X.device if isinstance(X, torch.Tensor) else None
+        generator = check_random_state(self.random_state, device)
         Xt = torch.as_tensor(X) if not isinstance(X, torch.Tensor) else X
         n_samples = int(Xt.shape[0])
         k = _resolve_n_samples(self.max_samples, n_samples)
@@ -131,15 +133,15 @@ def fit(self, X, y):
         members, _ = self._fit_members(X, y, bool(self.bootstrap))
         Xt = torch.as_tensor(X) if not isinstance(X, torch.Tensor) else X
         yt = torch.as_tensor(y) if not isinstance(y, torch.Tensor) else y
-        flat = yt.reshape(-1)
+        flat = yt.reshape(-1).to(Xt.device)
         try:
             uniq = sorted(set(flat.tolist()))
         except TypeError as e:
             raise TypeError("Labels must be sortable.") from e
         numeric = all(
             isinstance(v, (int, float)) and not isinstance(v, bool) for v in uniq
         )
-        self.classes_ = torch.as_tensor(uniq) if numeric else uniq
+        self.classes_ = torch.as_tensor(uniq, device=Xt.device) if numeric else uniq
         self.estimators_ = [clone(base).fit(Xt[idx], flat[idx]) for _, idx in members]
         self.n_features_in_ = int(Xt.shape[1]) if Xt.ndim == 2 else 0
         return self
@@ -219,7 +221,7 @@ def predict(self, X):
                     device=device,
                 )
             )
-        idx = torch.mode(torch.stack(cols), dim=0).values
+            idx = _majority_vote(torch.stack(cols))
         if isinstance(self.classes_, torch.Tensor):
             return self.classes_[idx]
         return [self.classes_[int(i)] for i in idx.tolist()]
@@ -296,7 +298,7 @@ def fit(self, X, y):
             if not isinstance(y, torch.Tensor)
             else y.to(dtype=_dtype)
         )
-        flat = yt.reshape(-1)
+        flat = yt.reshape(-1).to(Xt.device)
         self.estimators_ = [clone(base).fit(Xt[idx], flat[idx]) for _, idx in members]
         self.n_features_in_ = int(Xt.shape[1]) if Xt.ndim == 2 else 0
         return self

torml/ensemble/_forest.py

@@ -5,6 +5,7 @@
 import torch
 
 from torml.base import ClassifierMixin, RegressorMixin
+from torml.ensemble._voting import _majority_vote
 from torml.tree import DecisionTreeClassifier, DecisionTreeRegressor
 from torml.utils._random import check_random_state
 from torml.utils._validation import check_is_fitted
@@ -61,10 +62,11 @@ def fit(self, X, y):
             )
         if int(self.n_estimators) < 1:
             raise ValueError(f"n_estimators must be >= 1, got {self.n_estimators}.")
-        generator = check_random_state(self.random_state)
+        device = X.device if isinstance(X, torch.Tensor) else None
+        generator = check_random_state(self.random_state, device)
         Xt = torch.as_tensor(X) if not isinstance(X, torch.Tensor) else X
         yt = torch.as_tensor(y) if not isinstance(y, torch.Tensor) else y
-        flat = yt.reshape(-1)
+        flat = yt.reshape(-1).to(Xt.device)
         n_samples = int(Xt.shape[0])
         try:
             uniq = sorted(set(flat.tolist()))
@@ -73,7 +75,7 @@ def fit(self, X, y):
         numeric = all(
             isinstance(v, (int, float)) and not isinstance(v, bool) for v in uniq
         )
-        self.classes_ = torch.as_tensor(uniq) if numeric else uniq
+        self.classes_ = torch.as_tensor(uniq, device=Xt.device) if numeric else uniq
         self.estimators_ = []
         for _ in range(int(self.n_estimators)):
             idx = torch.randint(
@@ -120,7 +122,7 @@ def predict(self, X):
                     device=device,
                 )
             )
-        idx = torch.mode(torch.stack(cols), dim=0).values
+            idx = _majority_vote(torch.stack(cols))
         if isinstance(self.classes_, torch.Tensor):
             return self.classes_[idx]
         return [self.classes_[int(i)] for i in idx.tolist()]
@@ -175,15 +177,16 @@ def fit(self, X, y):
             )
         if int(self.n_estimators) < 1:
             raise ValueError(f"n_estimators must be >= 1, got {self.n_estimators}.")
-        generator = check_random_state(self.random_state)
+        device = X.device if isinstance(X, torch.Tensor) else None
+        generator = check_random_state(self.random_state, device)
         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=_dtype)
             if not isinstance(y, torch.Tensor)
             else y.to(dtype=_dtype)
         )
-        flat = yt.reshape(-1)
+        flat = yt.reshape(-1).to(Xt.device)
         n_samples = int(Xt.shape[0])
         self.estimators_ = []
         for _ in range(int(self.n_estimators)):

torml/ensemble/_voting.py

@@ -7,6 +7,33 @@
 from torml.base import BaseEstimator, ClassifierMixin, RegressorMixin, clone
 
 
+def _majority_vote(cols: torch.Tensor) -> torch.Tensor:
+    """Row-wise majority vote (smallest index wins ties, like mode).
+
+    Parameters
+    ----------
+    cols : torch.Tensor of shape (n_voters, n_queries)
+        Integer class indices per voter.
+
+    Returns
+    -------
+    winners : torch.Tensor of shape (n_queries,)
+        Winning class index per query.
+    """
+    counts = torch.zeros(
+        int(cols.shape[1]),
+        int(cols.max()) + 1,
+        dtype=torch.float32,
+        device=cols.device,
+    )
+    counts.scatter_add_(
+        1,
+        cols.T,
+        torch.ones_like(cols.T, dtype=torch.float32),
+    )
+    return torch.argmax(counts, dim=1)
+
+
 def _validate_estimators(estimators) -> list:
     """Validate the ``estimators`` list of (name, estimator) pairs."""
     if not isinstance(estimators, (list, tuple)) or len(estimators) == 0:
@@ -130,7 +157,9 @@ def fit(self, X, y):
         numeric = all(
             isinstance(v, (int, float)) and not isinstance(v, bool) for v in uniq
         )
-        self.classes_ = torch.as_tensor(uniq) if numeric else uniq
+        device = X.device if isinstance(X, torch.Tensor) else None
+        kwargs = {"device": device} if device is not None else {}
+        self.classes_ = torch.as_tensor(uniq, **kwargs) if numeric else uniq
         self.estimators_ = [(name, clone(est).fit(X, y)) for name, est in members]
         first = self.estimators_[0][1]
         if hasattr(first, "n_features_in_"):
@@ -232,7 +261,7 @@ def predict(self, X):
                         device=device,
                     )
                 )
-            idx = torch.mode(torch.stack(cols), dim=0).values
+            idx = _majority_vote(torch.stack(cols))
         if isinstance(self.classes_, torch.Tensor):
             return self.classes_[idx]
         return [self.classes_[int(i)] for i in idx.tolist()]

torml/feature_selection/_univariate.py

@@ -103,7 +103,7 @@ def f_classif(X, y):
     """
     Xt, yt = check_X_y(X, y)
     Xt = Xt.to(dtype=Xt.dtype)
-    flat = yt.reshape(-1)
+    flat = yt.reshape(-1).to(Xt.device)
     classes = torch.unique(flat, sorted=True)
     n = int(Xt.shape[0])
     n_classes = int(classes.shape[0])

torml/linear_model.py

@@ -88,7 +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)
-        y = y.to(dtype=X.dtype).reshape(-1, 1)
+        y = y.to(dtype=X.dtype, device=X.device).reshape(-1, 1)
         self.n_features_in_ = X.shape[1]
 
         if self.fit_intercept:
@@ -268,8 +268,8 @@ def fit(self, X: torch.Tensor, y: torch.Tensor):
                 "LogisticRegression supports binary labels only, "
                 f"got {int(classes.shape[0])} classes."
             )
-        self.classes_ = classes
-        target = inverse.to(dtype=X.dtype)
+        self.classes_ = classes.to(device=X.device)
+        target = inverse.to(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 * float(self.C))

torml/mixture/_gmm.py

@@ -97,7 +97,8 @@ 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)
+        device = X.device if isinstance(X, torch.Tensor) else None
+        generator = check_random_state(self.random_state, device)
         Xt = check_array(X, ensure_2d=True)
         n, d = int(Xt.shape[0]), int(Xt.shape[1])
         if n < int(self.n_components):

torml/model_selection/_split.py

@@ -81,7 +81,10 @@ def train_test_split(
     if not isinstance(shuffle, bool):
         raise TypeError(f"shuffle must be a bool, got {type(shuffle).__name__}.")
 
-    generator = check_random_state(random_state)
+    generator = check_random_state(
+        random_state,
+        arrays[0].device if isinstance(arrays[0], torch.Tensor) else None,
+    )
 
     n_samples = _num_samples(arrays[0])
     if n_samples < 2:
@@ -233,7 +236,9 @@ def split(self, X, y=None, groups=None):
             )
         if not self.shuffle and self.random_state is not None:
             raise ValueError("random_state must be None when shuffle=False.")
-        generator = check_random_state(self.random_state)
+        generator = check_random_state(
+            self.random_state, X.device if isinstance(X, torch.Tensor) else None
+        )
 
         n_samples = _num_samples(X)
         if n_samples < self.n_splits:

torml/multiclass/_ovr.py

@@ -73,7 +73,7 @@ def fit(self, X, y):
         numeric = all(
             isinstance(v, (int, float)) and not isinstance(v, bool) for v in uniq
         )
-        self.classes_ = torch.as_tensor(uniq) if numeric else uniq
+        self.classes_ = torch.as_tensor(uniq, device=Xt.device) if numeric else uniq
         self.estimators_ = []
         for c in uniq:
             cond = torch.as_tensor([v == c for v in flat.tolist()], device=Xt.device)

torml/neighbors/_classification.py

@@ -74,8 +74,9 @@ def fit(self, X, y):
             )
         flat_y = y.reshape(-1)
         self.classes_, inverse = torch.unique(flat_y, sorted=True, return_inverse=True)
+        self.classes_ = self.classes_.to(device=X.device)
         self.X_ = X
-        self.y_ = inverse.to(dtype=torch.long)
+        self.y_ = inverse.to(dtype=torch.long, device=X.device)
         self.n_features_in_ = int(X.shape[1])
         self.n_samples_in_ = int(X.shape[0])
         return self

torml/neighbors/_regression.py

@@ -65,7 +65,7 @@ def fit(self, X, y):
         p = _validate_p(self.p)
         self._p = p
         X, y = check_X_y(X, y)
-        y = y.to(dtype=X.dtype).reshape(-1)
+        y = y.to(dtype=X.dtype, device=X.device).reshape(-1)
         if int(X.shape[0]) < int(self.n_neighbors):
             raise ValueError(
                 f"n_neighbors ({self.n_neighbors}) must be <= n_samples "

torml/preprocessing/_label.py

@@ -88,7 +88,9 @@ def transform(self, y):
             if key not in self._class_to_index:
                 raise ValueError(f"Unknown label {v!r} seen in transform.")
             out.append(self._class_to_index[key])
-        return torch.tensor(out, dtype=torch.long)
+        device = y.device if isinstance(y, torch.Tensor) else None
+        kwargs = {"device": device} if device is not None else {}
+        return torch.tensor(out, dtype=torch.long, **kwargs)
 
     def fit_transform(self, y):
         """Fit and encode ``y``.
@@ -137,5 +139,7 @@ def inverse_transform(self, y):
                 )
             decoded.append(self._classes_list[ii])
         if isinstance(self.classes_, torch.Tensor):
-            return torch.as_tensor(decoded, dtype=self.classes_.dtype)
+            device = y.device if isinstance(y, torch.Tensor) else None
+            kwargs = {"device": device} if device is not None else {}
+            return torch.as_tensor(decoded, dtype=self.classes_.dtype, **kwargs)
         return decoded

torml/random_projection/_gaussian.py

@@ -81,7 +81,8 @@ def fit(self, X, y=None):
         self : GaussianRandomProjection
             Fitted projector.
         """
-        generator = check_random_state(self.random_state)
+        device = X.device if isinstance(X, torch.Tensor) else None
+        generator = check_random_state(self.random_state, device)
         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