wrap sklearn scalars for torch

torch_scalar.py

"""
how to wrap a scikit-learn scalar like RobustScaler for pytorch
"""
import torch
import numpy as np
from einops import rearrange
from sklearn.preprocessing import StandardScaler, RobustScaler

class TorchRobustScaler(RobustScaler):

    def wrap(self, X, method: str):
        b, l, h, v = X.shape
        X = rearrange(X, "b l h v -> b (l h v)")
        X = getattr(super(), method)(X)
        if isinstance(X, np.ndarray):
            X = torch.from_numpy(rearrange(X, "b (l h v) -> b l h v", l=l, h=h, v=v))
        return X


    def fit(self, X):
        return self.wrap(X, "fit")

    def transform(self, X):
        return self.wrap(X, "transform")

    def inverse_transform(self, X):
        return self.wrap(X, "inverse_transform")