|
5 | 5 | # SPDX-FileCopyrightText: Eric Martin <eric@ericmart.in> |
6 | 6 | # SPDX-FileCopyrightText: Giorgio Patrini <giorgio.patrini@anu.edu.au> |
7 | 7 | # SPDX-FileCopyrightText: Eric Chang <ericchang2017@u.northwestern.edu> |
8 | | -# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. |
| 8 | +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. |
9 | 9 | # SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause |
10 | 10 |
|
11 | 11 | # Original authors from Sckit-Learn: |
|
40 | 40 | SparseInputTagMixin, |
41 | 41 | StatelessTagMixin, |
42 | 42 | ) |
| 43 | +from cuml.internals.interop import InteropMixin, to_cpu, to_gpu |
43 | 44 |
|
44 | 45 | from ....common.array_descriptor import CumlArrayDescriptor |
45 | 46 | from ....internals.array import CumlArray |
@@ -519,6 +520,7 @@ def minmax_scale(X, feature_range=(0, 1), *, axis=0, copy=True): |
519 | 520 |
|
520 | 521 | class StandardScaler(TransformerMixin, |
521 | 522 | BaseEstimator, |
| 523 | + InteropMixin, |
522 | 524 | AllowNaNTagMixin, |
523 | 525 | SparseInputTagMixin): |
524 | 526 | """Standardize features by removing the mean and scaling to unit variance |
@@ -658,6 +660,47 @@ def _get_param_names(cls): |
658 | 660 | "copy" |
659 | 661 | ] |
660 | 662 |
|
| 663 | + # InteropMixin requirements |
| 664 | + _cpu_class_path = "sklearn.preprocessing.StandardScaler" |
| 665 | + |
| 666 | + @classmethod |
| 667 | + def _params_from_cpu(cls, model): |
| 668 | + """Convert sklearn StandardScaler hyperparameters to cuML format.""" |
| 669 | + return { |
| 670 | + "copy": model.copy, |
| 671 | + "with_mean": model.with_mean, |
| 672 | + "with_std": model.with_std, |
| 673 | + } |
| 674 | + |
| 675 | + def _params_to_cpu(self): |
| 676 | + """Convert cuML StandardScaler hyperparameters to sklearn format.""" |
| 677 | + return { |
| 678 | + "copy": self.copy, |
| 679 | + "with_mean": self.with_mean, |
| 680 | + "with_std": self.with_std, |
| 681 | + } |
| 682 | + |
| 683 | + def _attrs_from_cpu(self, model): |
| 684 | + """Convert sklearn StandardScaler fitted attributes to cuML format.""" |
| 685 | + attrs = { |
| 686 | + "mean_": to_gpu(mean) if (mean := getattr(model, "mean_", None)) is not None else None, |
| 687 | + "var_": to_gpu(var) if (var := getattr(model, "var_", None)) is not None else None, |
| 688 | + "scale_": to_gpu(scale) if (scale := getattr(model, "scale_", None)) is not None else None, |
| 689 | + "n_samples_seen_": to_gpu(nss) if (nss := getattr(model, "n_samples_seen_", None)) is not None else None, |
| 690 | + } |
| 691 | + return {**attrs, **super()._attrs_from_cpu(model)} |
| 692 | + |
| 693 | + def _attrs_to_cpu(self, model): |
| 694 | + """Convert cuML StandardScaler fitted attributes to sklearn format.""" |
| 695 | + |
| 696 | + attrs = { |
| 697 | + "mean_": to_cpu(mean) if (mean := getattr(self, "mean_", None)) is not None else None, |
| 698 | + "var_": to_cpu(var) if (var := getattr(self, "var_", None)) is not None else None, |
| 699 | + "scale_": to_cpu(scale) if (scale := getattr(self, "scale_", None)) is not None else None, |
| 700 | + "n_samples_seen_": None if (nss := getattr(self, "n_samples_seen_", None)) is None else cpu_np.int64(nss) if cpu_np.isscalar(nss) else to_cpu(nss), |
| 701 | + } |
| 702 | + return {**attrs, **super()._attrs_to_cpu(model)} |
| 703 | + |
661 | 704 | @reflect(reset=True) |
662 | 705 | def fit(self, X, y=None) -> "StandardScaler": |
663 | 706 | """Compute the mean and std to be used for later scaling. |
|
0 commit comments