autogluon.tabular 1.3.2b20250713__py3-none-any.whl → 1.3.2b20250714__py3-none-any.whl

autogluon/tabular/models/__init__.py

@@ -23,6 +23,7 @@ from .tabicl.tabicl_model import TabICLModel
 from .tabm.tabm_model import TabMModel
 from .tabpfnv2.tabpfnv2_model import TabPFNV2Model
 from .tabpfnmix.tabpfnmix_model import TabPFNMixModel
+from .mitra.mitra_model import MitraModel
 from .tabular_nn.torch.tabular_nn_torch import TabularNeuralNetTorchModel
 from .text_prediction.text_prediction_v1_model import TextPredictorModel
 from .xgboost.xgboost_model import XGBoostModel

autogluon/tabular/models/mitra/_internal/config/config_pretrain.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+import yaml
+import os
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+
+from ..._internal.config.enums import GeneratorName, ModelName, LossName, Task
+
+@dataclass
+class ConfigData():
+    generator: GeneratorName
+    min_samples_support: int
+    max_samples_support: int
+    n_samples_query: int
+    min_features: int
+    max_features: int
+    max_classes: int
+    sample_multinomial_categorical: bool
+    sample_multinomial_label: bool
+    generator_hyperparams: dict
+    task: Task
+
+    def __post_init__(self):
+
+        assert self.min_samples_support <= self.max_samples_support
+        assert self.min_features <= self.max_features
+
+@dataclass
+class ConfigModel():
+    name: ModelName
+    hyperparams: dict
+
+
+@dataclass
+class ConfigPreprocessing():
+    use_quantile_transformer: bool
+    use_feature_count_scaling: bool
+
+@dataclass
+class ConfigGradScaler():
+    enabled: bool
+    scale_init: float
+    scale_min: float
+    growth_interval: int
+
+
+    def __post_init__(self):
+        assert self.scale_init >= self.scale_min, "Scale init must be greater than scale min"
+        assert self.scale_min >= 1, "Scale min lower than 1 makes no sense for mixed precision training"
+        assert type(self.scale_init) == float, "Scale init must be a float, otherwise gradscaler will return an error"
+        assert type(self.scale_min) == float, "Scale min must be a float, otherwise gradscaler will return an error"
+
+@dataclass
+class ConfigOptim():
+    steps: int
+    log_every_n_steps: int
+    eval_every_n_steps: int
+    batch_size: int
+    gradient_accumulation_steps: int
+    lr: float
+    weight_decay: float
+    beta1: float
+    beta2: float
+    warmup_steps: int
+    cosine_scheduler: bool
+    max_grad_norm: float
+    label_smoothing: float
+    regression_loss: LossName
+    use_pretrained_weights: bool
+    path_to_weights: str
+    resume_states: bool
+    path_to_states: str
+    precision: str
+    grad_scaler: ConfigGradScaler
+
+    @classmethod
+    def from_hydra(cls, cfg_hydra: DictConfig) -> Self:
+
+        grad_scaler = ConfigGradScaler(**cfg_hydra.grad_scaler)
+        cfg_dict: dict = OmegaConf.to_container(cfg_hydra)      # type: ignore
+        del cfg_dict["grad_scaler"]
+
+        regression_loss = LossName[cfg_dict["regression_loss"]]
+        del cfg_dict["regression_loss"]
+
+        return cls(
+            grad_scaler=grad_scaler,
+            regression_loss=regression_loss,
+            **cfg_dict
+        )
+
+    def __post_init__(self):
+        assert hasattr(torch, self.precision), f"Precision {self.precision} not supported by torch"
+
+class ConfigSaveLoadMixin(yaml.YAMLObject):
+
+    def save(self, path: Path) -> None:
+
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(path, 'w') as f:
+            yaml.dump(self, f, default_flow_style=False)
+
+
+    @classmethod
+    def load(cls, path: Path) -> Self:
+
+        with open(path, 'r') as f:
+            # It's unsafe, but not unsafer than the pickle module
+            config = yaml.unsafe_load(f)
+
+        return config
+
+@dataclass
+class ConfigPretrain(ConfigSaveLoadMixin):
+    run_name: str
+    output_dir: Path
+    seed: int
+    devices: list[torch.device]
+    device: torch.device
+    max_cpus_per_device: Optional[int]
+    use_ddp: bool
+    workers_per_gpu: int
+    model: ConfigModel
+    data: ConfigData
+    optim: ConfigOptim
+    preprocessing: ConfigPreprocessing
+    load_from_file: bool
+    load_path_x: str
+    load_path_y: str
+    save_file: bool
+    save_file_only: bool
+    save_path_x: str
+    save_path_y: str
+    number_of_runs: int
+
+    @classmethod
+    def from_hydra(cls, cfg_hydra: DictConfig):
+
+        assert not os.path.exists(cfg_hydra.output_dir), f'Output directory {cfg_hydra.output_dir} already exists! Please change to a new folder.'
+
+        output_dir = Path(cfg_hydra.output_dir)
+
+        devices = [torch.device(device) for device in cfg_hydra.devices]
+
+        # Initialize device to cpu, DDP will overwrite this
+        device = torch.device("cpu")
+
+        return cls(
+            run_name=cfg_hydra.run_name,
+            output_dir=output_dir,
+            devices=devices,
+            device=device,
+            max_cpus_per_device=cfg_hydra.max_cpus_per_device,
+            use_ddp=len(devices) > 1,
+            seed=cfg_hydra.seed,
+            workers_per_gpu=cfg_hydra.workers_per_gpu,
+            model = ConfigModel(
+                name = ModelName[cfg_hydra.model.name],
+                hyperparams = OmegaConf.to_container(cfg_hydra.model.hyperparams),   
+            ),
+            data = ConfigData(
+                generator=GeneratorName(cfg_hydra.data.generator),
+                min_samples_support=cfg_hydra.data.min_samples_support,
+                max_samples_support=cfg_hydra.data.max_samples_support,
+                n_samples_query=cfg_hydra.data.n_samples_query,
+                min_features=cfg_hydra.data.min_features,
+                max_features=cfg_hydra.data.max_features,
+                max_classes=cfg_hydra.data.max_classes,
+                task=Task[cfg_hydra.data.task],
+                sample_multinomial_categorical=cfg_hydra.data.sample_multinomial_categorical,
+                sample_multinomial_label=cfg_hydra.data.sample_multinomial_label,
+                generator_hyperparams=OmegaConf.to_container(cfg_hydra.data.generator_hyperparams),    # type: ignore
+            ),
+            optim = ConfigOptim.from_hydra(cfg_hydra.optim),
+            preprocessing = ConfigPreprocessing(**cfg_hydra.preprocessing),
+            load_from_file = cfg_hydra.load_from_file,
+            load_path_x = cfg_hydra.load_path_x,
+            load_path_y = cfg_hydra.load_path_y,
+            save_file = cfg_hydra.save_file,
+            save_file_only = cfg_hydra.save_file_only,
+            save_path_x = cfg_hydra.save_path_x,
+            save_path_y = cfg_hydra.save_path_y,
+            number_of_runs = cfg_hydra.number_of_runs,
+        )

autogluon/tabular/models/mitra/_internal/config/config_run.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Self
+
+import torch
+
+from ..._internal.config.config_pretrain import ConfigSaveLoadMixin
+from ..._internal.config.enums import ModelName
+
+@dataclass
+class ConfigRun(ConfigSaveLoadMixin):
+    device: torch.device
+    seed: int
+    model_name: ModelName
+    hyperparams: dict
+
+    @classmethod
+    def create(
+        cls,  
+        device: torch.device, 
+        seed: int, 
+        model_name: ModelName, 
+        hyperparams: dict
+    ) -> Self:
+       
+        return cls(
+            device=device,
+            seed=seed,
+            model_name=model_name,
+            hyperparams=hyperparams
+    )

autogluon/tabular/models/mitra/_internal/config/enums.py

@@ -0,0 +1,145 @@
+from enum import IntEnum, StrEnum
+
+
+class Task(StrEnum):
+    CLASSIFICATION = "classification"
+    REGRESSION = "regression"
+
+
+class FeatureType(StrEnum):
+    NUMERICAL = "numerical"
+    CATEGORICAL = "categorical"
+    MIXED = "mixed"
+
+
+class SearchType(StrEnum):
+    DEFAULT = "default"
+    RANDOM = "random"
+
+
+class DatasetSize(IntEnum):
+    SMALL = 1000
+    MEDIUM = 10000
+    LARGE = 50000
+
+
+class DataSplit(StrEnum):
+    TRAIN = "train"
+    VALID = "valid"
+    TEST = "test"
+
+
+class Phase(StrEnum):
+    TRAINING = "training"
+    VALIDATION = "validation"
+    TESTING = "testing"
+
+
+class ModelName(StrEnum):
+    PLACEHOLDER = "_placeholder_"   # This is a placeholder for the current running model
+    FT_TRANSFORMER = "FT-Transformer"
+    TABPFN = "TabPFN"
+    FOUNDATION = "Foundation"
+    FOUNDATION_FLASH = "FoundationFlash"
+    TAB2D = "Tab2D"
+    TAB2D_COL_ROW = "Tab2D_COL_ROW"
+    TAB2D_SDPA = "Tab2D_SDPA"
+    SAINT = "SAINT"
+    MLP = "MLP"
+    MLP_RTDL = "MLP-rtdl"
+    RESNET = "Resnet"
+    RANDOM_FOREST = "RandomForest"
+    XGBOOST = "XGBoost"
+    CATBOOST = "CatBoost"
+    LIGHTGBM = "LightGBM"
+    GRADIENT_BOOSTING_TREE = "GradientBoostingTree"
+    HIST_GRADIENT_BOOSTING_TREE = "HistGradientBoostingTree"
+    LOGISTIC_REGRESSION = "LogisticRegression"
+    LINEAR_REGRESSION = "LinearRegression"
+    DECISION_TREE = "DecisionTree"
+    KNN = "KNN"
+    STG = "STG"
+    SVM = "SVM"
+    TABNET = "TabNet"
+    TABTRANSFORMER = "TabTransformer"
+    DEEPFM = "DeepFM"
+    VIME = "VIME"
+    DANET = "DANet"
+    NODE = "NODE"
+    AUTOGLUON = "AutoGluon"
+
+
+class ModelClass(StrEnum):
+    BASE = 'base'
+    GBDT = 'GBDT'
+    NN = 'NN'
+    ICLT = 'ICLT'
+
+
+class DownstreamTask(StrEnum):
+    ZEROSHOT = "zeroshot"
+    FINETUNE = "finetune"
+
+
+
+class BenchmarkName(StrEnum):
+    DEBUG_CLASSIFICATION = "debug_classification"
+    DEBUG_REGRESSION = "debug_regression"
+    DEBUG_TABZILLA = "debug_tabzilla"
+
+    CATEGORICAL_CLASSIFICATION = "categorical_classification"
+    NUMERICAL_CLASSIFICATION = "numerical_classification"
+    CATEGORICAL_REGRESSION = "categorical_regression"
+    NUMERICAL_REGRESSION = "numerical_regression"
+    CATEGORICAL_CLASSIFICATION_LARGE = "categorical_classification_large"
+    NUMERICAL_CLASSIFICATION_LARGE = "numerical_classification_large"
+    CATEGORICAL_REGRESSION_LARGE = "categorical_regression_large"
+    NUMERICAL_REGRESSION_LARGE = "numerical_regression_large"
+    
+    TABZILLA_HARD = "tabzilla_hard"
+    TABZILLA_HARD_MAX_TEN_CLASSES = "tabzilla_hard_max_ten_classes"
+    TABZILLA_HAS_COMPLETED_RUNS = "tabzilla_has_completed_runs"
+
+
+class BenchmarkOrigin(StrEnum):
+    TABZILLA = "tabzilla"
+    WHYTREES = "whytrees"
+
+
+class GeneratorName(StrEnum):
+    TABPFN = 'tabpfn'
+    TREE = 'tree'
+    RANDOMFOREST = 'randomforest'
+    NEIGHBOR = 'neighbor'
+    MIX = 'mix'
+    PERLIN = 'perlin'
+    MIX_7 = 'mix_7'
+    MIX_6 = 'mix_6'
+    MIX_5 = 'mix_5'
+    MIX_5_GP = 'mix_5_gp'
+    MIX_4 = 'mix_4'
+    MIX_4_AG = 'mix_4_ag'
+    LR = 'lr'
+    POLY = 'poly'
+    SAMPLE_RF = 'sample_rf'
+    SAMPLE_GP = 'sample_gp'
+    TABREPO = 'tabrepo'
+    MIX_4_TABREPO = 'mix_4_tabrepo'
+    MIX_4_TABPFNV2 = 'mix_4_tabpfnv2'
+
+
+class MetricName(StrEnum):
+    ACCURACY = "accuracy"
+    F1 = "f1"
+    AUC = "auc"
+    MSE = "mse"
+    MAE = "mae"
+    R2 = "r2"
+    LOG_LOSS = "log_loss"
+    RMSE = "rmse"
+
+
+class LossName(StrEnum):
+    CROSS_ENTROPY = "cross_entropy"
+    MSE = "mse"
+    MAE = "mae"

autogluon/tabular/models/mitra/_internal/core/callbacks.py

@@ -0,0 +1,94 @@
+import numpy as np
+import torch
+
+
+class EarlyStopping():
+
+    def __init__(self, patience=10, delta=0.0001, metric='log_loss'):
+
+        self.patience = patience
+        self.counter = 0
+        self.best_score = None
+        self.early_stop = False
+        self.delta = delta
+        self.metric = metric
+
+
+    def __call__(self, val_loss):
+        
+        # smaller is better for these metrics
+        if self.metric in ["log_loss", "mse", "mae", "rmse"]:
+            score = -val_loss
+        # larger is better for these metrics
+        elif self.metric in ["accuracy", "roc_auc", "r2"]:
+            score = val_loss
+        else:
+            raise ValueError(f"Unsupported metric: {self.metric}. Supported metrics are: log_loss, mse, mae, rmse, accuracy, roc_auc, r2.")
+
+        if self.best_score is None:
+            self.best_score = score
+        elif score < self.best_score + self.delta:
+            self.counter += 1
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_score = score
+            self.counter = 0
+
+    def we_should_stop(self):
+        return self.early_stop
+
+
+class Checkpoint():
+
+    def __init__(self):
+        self.curr_best_loss = np.inf
+        self.best_model: dict
+    
+    def reset(self, net: torch.nn.Module):
+        self.curr_best_loss = np.inf
+        self.best_model = net.state_dict()
+        for key in self.best_model:
+            self.best_model[key] = self.best_model[key].to('cpu')
+        
+
+    def __call__(self, net: torch.nn.Module, loss: float):
+        
+        if loss < self.curr_best_loss:
+            self.curr_best_loss = loss
+            self.best_model = net.state_dict()
+            for key in self.best_model:
+                self.best_model[key] = self.best_model[key].to('cpu')
+
+
+    def set_to_best(self, net):
+        net.load_state_dict(self.best_model)
+
+
+class EpochStatistics():
+
+    def __init__(self) -> None:
+        self.n = 0
+        self.loss = 0
+        self.score = 0
+        
+    def update(self, loss, score, n):
+        self.n += n
+        self.loss += loss * n
+        self.score += score * n
+
+    def get(self):
+        return self.loss / self.n, self.score / self.n
+    
+class TrackOutput():
+
+    def __init__(self) -> None:
+        self.y_true: list[np.ndarray] = []
+        self.y_pred: list[np.ndarray] = []
+
+    def update(self, y_true: np.ndarray, y_pred: np.ndarray):
+        self.y_true.append(y_true)
+        self.y_pred.append(y_pred)
+
+    def get(self):
+        return np.concatenate(self.y_true, axis=0), np.concatenate(self.y_pred, axis=0)

autogluon/tabular/models/mitra/_internal/core/get_loss.py

@@ -0,0 +1,55 @@
+import torch
+import einops
+
+from ..._internal.config.config_pretrain import ConfigPretrain
+from ..._internal.config.config_run import ConfigRun
+from ..._internal.config.enums import LossName, Task
+
+class CrossEntropyLossExtraBatch(torch.nn.Module):
+
+    def __init__(self, label_smoothing: float):
+        super().__init__()
+
+        self.loss = torch.nn.CrossEntropyLoss(label_smoothing=label_smoothing)
+
+
+    def forward(self, input, target):
+        """
+        Input has shape (batch_size, num_samples, num_classes)
+        Target has shape (batch_size, num_samples)
+
+        Compared to the original CrossEntropyLoss, accepts (batch_size, num_samples) as batch
+        """
+
+        input = einops.rearrange(input, 'b s c -> (b s) c')
+        target = einops.rearrange(target, 'b s -> (b s)')
+
+        return self.loss(input, target)
+
+def get_loss(cfg: ConfigRun):
+
+    match (cfg.task, cfg.hyperparams['regression_loss']):
+        case (Task.REGRESSION, LossName.MSE):
+            return torch.nn.MSELoss()
+        case (Task.REGRESSION, LossName.MAE):
+            return torch.nn.L1Loss()
+        case (Task.REGRESSION, LossName.CROSS_ENTROPY):
+            return CrossEntropyLossExtraBatch(cfg.hyperparams['label_smoothing'])
+        case (Task.CLASSIFICATION, _):
+            return CrossEntropyLossExtraBatch(cfg.hyperparams['label_smoothing'])
+        case (_, _):
+            raise ValueError(f"Unsupported task {cfg.task} and (regression) loss {cfg.hyperparams['regression_loss']}")
+
+def get_loss_pretrain(cfg: ConfigPretrain):
+
+    match (cfg.data.task, cfg.optim.regression_loss):
+        case (Task.REGRESSION, LossName.MSE):
+            return torch.nn.MSELoss()
+        case (Task.REGRESSION, LossName.MAE):
+            return torch.nn.L1Loss()
+        case (Task.REGRESSION, LossName.CROSS_ENTROPY):
+            return CrossEntropyLossExtraBatch(cfg.optim.label_smoothing)
+        case (Task.CLASSIFICATION, _):
+            return CrossEntropyLossExtraBatch(cfg.optim.label_smoothing)
+        case (_, _):
+            raise ValueError(f"Unsupported task {cfg.data.task} and (regression) loss {cfg.optim.regression_loss}")

autogluon/tabular/models/mitra/_internal/core/get_optimizer.py

@@ -0,0 +1,108 @@
+import torch
+from torch.optim import SGD, Adam, AdamW
+
+from ..._internal.config.config_pretrain import ConfigPretrain
+
+
+def get_optimizer(hyperparams: dict, model: torch.nn.Module) -> torch.optim.Optimizer:
+
+    optimizer: torch.optim.Optimizer
+
+    if hyperparams['optimizer'] == "adam":
+        optimizer = Adam(
+            model.parameters(), 
+            lr=hyperparams['lr'],
+            betas=(0.9, 0.999),
+            weight_decay=hyperparams['weight_decay']
+        )
+    elif hyperparams['optimizer'] == "adamw":
+        optimizer = AdamW(
+            model.parameters(), 
+            lr=hyperparams['lr'],
+            betas=(0.9, 0.999),
+            weight_decay=hyperparams['weight_decay']
+        )
+    elif hyperparams['optimizer'] == "sgd":
+        optimizer = SGD(
+            model.parameters(),
+            lr=hyperparams['lr'],
+            weight_decay=hyperparams['weight_decay']
+        )
+    else:
+        raise ValueError("Optimizer not recognized")
+    
+    return optimizer
+
+
+def get_optimizer_pretrain(cfg: ConfigPretrain, model: torch.nn.Module) -> torch.optim.Optimizer:
+
+    parameters = [(name, param) for name, param in model.named_parameters()]
+
+    parameters_with_weight_decay = []
+    parameters_without_weight_decay = []
+
+    for name, param in parameters:
+        if name.endswith("bias") or 'norm' in name or 'embedding' in name:
+            parameters_without_weight_decay.append(param)
+        else:
+            parameters_with_weight_decay.append(param)
+
+    optimizer_parameters = [
+        {"params": parameters_with_weight_decay, "weight_decay": cfg.optim.weight_decay},
+        {"params": parameters_without_weight_decay, "weight_decay": 0.0},
+    ]
+
+    optimizer = torch.optim.AdamW(
+        optimizer_parameters, 
+        lr=cfg.optim.lr,
+        betas=(cfg.optim.beta1, cfg.optim.beta2),
+        weight_decay=cfg.optim.weight_decay
+    )
+    
+    return optimizer
+
+
+class GradScaler(torch.amp.GradScaler):
+
+    def __init__(
+        self, 
+        enabled: bool = True,
+        scale_init: float = 2.**16,
+        scale_min: float = 1.,
+        growth_interval: int = 2000,
+        device: str = 'cuda'
+    ):
+        super().__init__(enabled=enabled, device="cpu", init_scale=scale_init, growth_interval=growth_interval) # type: ignore
+        self._enabled = enabled
+        self.scale_min = scale_min
+        self.device = device
+
+        if not self._enabled:
+            # We write scale=1 to log if the scaler is disabled
+            self._scale = torch.tensor((1,), dtype=torch.float32, device=self.device)
+
+
+    def update(self):
+
+        if not self._enabled:
+            return
+
+        super().update()
+
+        if self._scale < self.scale_min:
+            super().update(self.scale_min)
+
+
+def move_optimizer_to(optim, device):
+    for param in optim.state.values():
+        # Not sure there are any global tensors in the state dict
+        if isinstance(param, torch.Tensor):
+            param.data = param.data.to(device)
+            if param._grad is not None:
+                param._grad.data = param._grad.data.to(device)
+        elif isinstance(param, dict):
+            for subparam in param.values():
+                if isinstance(subparam, torch.Tensor):
+                    subparam.data = subparam.data.to(device)
+                    if subparam._grad is not None:
+                        subparam._grad.data = subparam._grad.data.to(device)

autogluon/tabular/models/mitra/_internal/core/get_scheduler.py

@@ -0,0 +1,67 @@
+import torch
+from torch.optim.lr_scheduler import ReduceLROnPlateau, LinearLR
+from transformers import get_constant_schedule_with_warmup
+from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup
+
+from ..._internal.config.config_pretrain import ConfigPretrain
+
+
+def get_scheduler(hyperparams: dict, optimizer: torch.optim.Optimizer) -> tuple[torch.optim.lr_scheduler.LambdaLR, ReduceLROnPlateau]:
+
+    warmup_steps = hyperparams['warmup_steps']
+
+    # if warmup_steps > 0:
+    #     scheduler_warmup = torch.optim.lr_scheduler.LambdaLR(
+    #         optimizer, lambda step: min((step + 1) / warmup_steps, 1.0)
+    #     )
+    # else:
+    #     scheduler_warmup = torch.optim.lr_scheduler.LambdaLR(
+    #         optimizer, lambda step: 1.0
+    #     )
+    
+    if warmup_steps > 0:
+        scheduler_warmup = LinearLR(
+            optimizer,
+            start_factor=1.0 / warmup_steps,
+            end_factor=1.0,
+            total_iters=warmup_steps,
+        )
+    else:
+        scheduler_warmup = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=1)
+
+    if hyperparams['lr_scheduler']:      
+        scheduler_reduce_on_plateau = ReduceLROnPlateau(
+            optimizer, 
+            patience=hyperparams['lr_scheduler_patience'], 
+            min_lr=0, 
+            factor=0.2
+        )
+    else:
+        # With ReduceLROnPlateau, the scheduler accepts a metric to monitor, so our dummy metric must also be a ReduceLRonPlateau scheduler
+        scheduler_reduce_on_plateau = ReduceLROnPlateau(
+            optimizer, 
+            patience=1000000000, 
+            min_lr=0, 
+            factor=0.2
+        )
+
+    return scheduler_warmup, scheduler_reduce_on_plateau
+
+
+def get_scheduler_pretrain(cfg: ConfigPretrain, optimizer: torch.optim.Optimizer):
+
+    
+    if cfg.optim.cosine_scheduler:
+        schedule = get_cosine_with_min_lr_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=cfg.optim.warmup_steps,
+            num_training_steps=cfg.optim.steps,
+            min_lr_rate=0.1
+        )
+    else:
+        schedule = get_constant_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=cfg.optim.warmup_steps
+        )
+
+    return schedule