autogluon.tabular 1.3.2b20250713__tar.gz → 1.3.2b20250715__tar.gz

{autogluon.tabular-1.3.2b20250713 → autogluon.tabular-1.3.2b20250715}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: autogluon.tabular
-Version: 1.3.2b20250713
+Version: 1.3.2b20250715
 Summary: Fast and Accurate ML in 3 Lines of Code
 Home-page: https://github.com/autogluon/autogluon
 Author: AutoGluon Community
@@ -42,6 +42,7 @@ Provides-Extra: fastai
 Provides-Extra: tabm
 Provides-Extra: tabpfn
 Provides-Extra: tabpfnmix
+Provides-Extra: mitra
 Provides-Extra: tabicl
 Provides-Extra: ray
 Provides-Extra: skex

{autogluon.tabular-1.3.2b20250713 → autogluon.tabular-1.3.2b20250715}/setup.py

@@ -62,6 +62,13 @@ extras_require = {
         "huggingface_hub[torch]",  # Only needed for HuggingFace downloads, currently uncapped to minimize future conflicts.
         "einops>=0.7,<0.9",
     ],
+    "mitra": [
+        "loguru",
+        "einx",
+        "omegaconf",
+        "transformers",
+        # "flash-attn>2.6.3,<2.8", # TODO: flash-attn installation requires --no-build-isolation and torch, python and cuda version compatibility.
+    ],
     "tabicl": [
         "tabicl>=0.1.3,<0.2",  # 0.1.3 added a major bug fix to multithreading.
     ],
@@ -109,7 +116,7 @@ extras_require["all"] = all_requires
 
 
 test_requires = []
-for test_package in ["tabpfn", "imodels", "skl2onnx", "tabicl"]:
+for test_package in ["tabpfn", "imodels", "skl2onnx", "tabicl", "mitra"]:
     test_requires += extras_require[test_package]
 extras_require["tests"] = test_requires
 install_requires = ag.get_dependency_version_ranges(install_requires)

{autogluon.tabular-1.3.2b20250713 → autogluon.tabular-1.3.2b20250715}/src/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-1.3.2b20250713 → autogluon.tabular-1.3.2b20250715}/src/autogluon/tabular/models/catboost/catboost_model.py

@@ -13,13 +13,13 @@ from autogluon.common.features.types import R_BOOL, R_CATEGORY, R_FLOAT, R_INT
 from autogluon.common.utils.pandas_utils import get_approximate_df_mem_usage
 from autogluon.common.utils.resource_utils import ResourceManager
 from autogluon.common.utils.try_import import try_import_catboost
-from autogluon.core.constants import MULTICLASS, PROBLEM_TYPES_CLASSIFICATION, QUANTILE, SOFTCLASS
+from autogluon.core.constants import MULTICLASS, PROBLEM_TYPES_CLASSIFICATION, REGRESSION, QUANTILE, SOFTCLASS
 from autogluon.core.models import AbstractModel
 from autogluon.core.models._utils import get_early_stopping_rounds
 from autogluon.core.utils.exceptions import TimeLimitExceeded
 
 from .callbacks import EarlyStoppingCallback, MemoryCheckCallback, TimeCheckCallback
-from .catboost_utils import get_catboost_metric_from_ag_metric
+from .catboost_utils import get_catboost_metric_from_ag_metric, CATBOOST_EVAL_METRIC_TO_LOSS_FUNCTION
 from .hyperparameters.parameters import get_param_baseline
 from .hyperparameters.searchspaces import get_default_searchspace
 
@@ -131,11 +131,14 @@ class CatBoostModel(AbstractModel):
             # FIXME: This is extremely slow due to unoptimized metric / objective sent to CatBoost
             from .catboost_softclass_utils import SoftclassCustomMetric, SoftclassObjective
 
-            params["loss_function"] = SoftclassObjective.SoftLogLossObjective()
+            params.setdefault("loss_function",  SoftclassObjective.SoftLogLossObjective())
             params["eval_metric"] = SoftclassCustomMetric.SoftLogLossMetric()
-        elif self.problem_type == QUANTILE:
-            # FIXME: Unless specified, CatBoost defaults to loss_function='MultiQuantile' and raises an exception
-            params["loss_function"] = params["eval_metric"]
+        elif self.problem_type in [REGRESSION, QUANTILE]:
+            # Choose appropriate loss_function that is as close as possible to the eval_metric
+            params.setdefault(
+                "loss_function",
+                CATBOOST_EVAL_METRIC_TO_LOSS_FUNCTION.get(params["eval_metric"], params["eval_metric"])
+            )
 
         model_type = CatBoostClassifier if self.problem_type in PROBLEM_TYPES_CLASSIFICATION else CatBoostRegressor
         num_rows_train = len(X)

{autogluon.tabular-1.3.2b20250713 → autogluon.tabular-1.3.2b20250715}/src/autogluon/tabular/models/catboost/catboost_utils.py

@@ -6,6 +6,13 @@ logger = logging.getLogger(__name__)
 
 
 CATBOOST_QUANTILE_PREFIX = "Quantile:"
+# Mapping from non-optimizable eval_metric to optimizable loss_function.
+# See https://catboost.ai/docs/en/concepts/loss-functions-regression#usage-information
+CATBOOST_EVAL_METRIC_TO_LOSS_FUNCTION = {
+    "MedianAbsoluteError": "MAE",
+    "SMAPE": "MAPE",
+    "R2": "RMSE",
+}
 
 
 # TODO: Add weight support?
@@ -65,7 +72,10 @@ def get_catboost_metric_from_ag_metric(metric, problem_type, quantile_levels=Non
             mean_squared_error="RMSE",
             root_mean_squared_error="RMSE",
             mean_absolute_error="MAE",
+            mean_absolute_percentage_error="MAPE",
+            # Non-optimizable metrics, see CATBOOST_EVAL_METRIC_TO_LOSS_FUNCTION
             median_absolute_error="MedianAbsoluteError",
+            symmetric_mean_absolute_percentage_error="SMAPE",
             r2="R2",
         )
         metric_class = metric_map.get(metric.name, "RMSE")

{autogluon.tabular-1.3.2b20250713 → autogluon.tabular-1.3.2b20250715}/src/autogluon/tabular/models/lgb/lgb_model.py

@@ -281,7 +281,8 @@ class LGBModel(AbstractModel):
                 train_params["params"]["metric"] = f'{stopping_metric},{train_params["params"]["metric"]}'
 
         if self.problem_type == SOFTCLASS:
-            train_params["fobj"] = lgb_utils.softclass_lgbobj
+            train_params["params"]["objective"] = lgb_utils.softclass_lgbobj
+            train_params["params"]["num_classes"] = self.num_classes
         elif self.problem_type == QUANTILE:
             train_params["params"]["quantile_levels"] = self.quantile_levels
         if seed_val is not None:

autogluon.tabular-1.3.2b20250715/src/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-1.3.2b20250715/src/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-1.3.2b20250715/src/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-1.3.2b20250715/src/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-1.3.2b20250715/src/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}")