Trainers Module

GeneralTrainer

Bases: BaseTrainer

Handles general training logic for autoencoder models.

Attributes:

Name	Type	Description
_trainset		The dataset used for training.
_validset		The dataset used for validation, if provided.
_result		An object to store and manage training results.
_config		Configuration object containing training hyperparameters and settings.
_model_type		The autoencoder model class to be trained.
_loss_fn		Instantiated loss function specific to the model.
_trainloader		DataLoader for the training dataset.
_validloader		DataLoader for the validation dataset, if provided.
_model		The instantiated model architecture.
_optimizer		The optimizer used for training.
_fabric		Lightning Fabric wrapper for device and precision management.
n_train		Number of training samples.
n_valid		Number of validation samples.
n_test	Optional[int]	Number of test samples (set during prediction).
n_features		Number of input features.
latent_dim		Dimensionality of the latent space.
device		Device on which the model is located.
_n_cpus		Number of CPU cores available.
_reconstruction_buffer		Buffer to store reconstructions during training/validation/testing.
_latentspace_buffer		Buffer to store latent representations during training/validation/testing.
_mu_buffer		Buffer to store latent means (for VAE) during training/validation
_sigma_buffer		Buffer to store latent log-variances (for VAE) during training/validation/testing.
_sample_ids_buffer		Buffer to store sample IDs during training/validation/testing.

Source code in src/autoencodix/trainers/_general_trainer.py

class GeneralTrainer(BaseTrainer):
    """Handles general training logic for autoencoder models.

    Attributes:
         _trainset: The dataset used for training.
         _validset: The dataset used for validation, if provided.
         _result: An object to store and manage training results.
         _config: Configuration object containing training hyperparameters and settings.
         _model_type: The autoencoder model class to be trained.
         _loss_fn: Instantiated loss function specific to the model.
         _trainloader: DataLoader for the training dataset.
         _validloader: DataLoader for the validation dataset, if provided.
         _model: The instantiated model architecture.
         _optimizer: The optimizer used for training.
         _fabric: Lightning Fabric wrapper for device and precision management.
         n_train: Number of training samples.
         n_valid: Number of validation samples.
         n_test: Number of test samples (set during prediction).
         n_features: Number of input features.
         latent_dim: Dimensionality of the latent space.
         device: Device on which the model is located.
         _n_cpus: Number of CPU cores available.
         _reconstruction_buffer: Buffer to store reconstructions during training/validation/testing.
         _latentspace_buffer: Buffer to store latent representations during training/validation/testing.
         _mu_buffer: Buffer to store latent means (for VAE) during training/validation
         _sigma_buffer: Buffer to store latent log-variances (for VAE) during training/validation/testing.
         _sample_ids_buffer: Buffer to store sample IDs during training/validation/testing.

    """

    def __init__(
        self,
        trainset: Optional[BaseDataset],
        validset: Optional[BaseDataset],
        result: Result,
        config: DefaultConfig,
        model_type: Type[BaseAutoencoder],
        loss_type: Type[BaseLoss],
        ontologies: Optional[Union[Tuple, List]] = None,
        **kwargs,
    ):
        """Initializes the GeneralTrainer.

        Args:
                trainset: The dataset used for training.
                validset: The dataset used for validation, if provided.
                result: An object to store and manage training results.
                config: Configuration object containing training hyperparameters and settings.
                model_type: The autoencoder model class to be trained.
                loss_type: The loss function class to be used for training.
                ontologies: Ontology information, if provided for Ontix
        """
        self._n_cpus = os.cpu_count()
        if self._n_cpus is None:
            self._n_cpus = 0

        super().__init__(
            trainset, validset, result, config, model_type, loss_type, ontologies
        )

        self.latent_dim = config.latent_dim
        if ontologies is not None:
            if not hasattr(self._model, "latent_dim"):
                raise ValueError(
                    "Model must have a 'latent_dim' attribute when ontologies are provided."
                )
            self.latent_dim = self._model.latent_dim

        # we will set this later, in the predict method
        self.n_test: Optional[int] = None
        self.n_train = len(trainset) if trainset else 0
        self.n_valid = len(validset) if validset else 0
        self.n_features = trainset.get_input_dim() if trainset else 0
        self.device = next(self._model.parameters()).device
        self._init_buffers()

    def _init_buffers(self, input_data: Optional[BaseDataset] = None) -> None:
        if input_data:
            self.n_features = input_data.get_input_dim()

        def make_tensor_buffer(size: int, dim: Union[int, Tuple[int, ...]]):
            if isinstance(dim, int):
                if self._config.save_vram:
                    return torch.zeros(
                        (size, dim), device="cpu"
                    )  # Move to CPU to avoid GPU OOM, but will be slower
                else:
                    return torch.zeros((size, dim), device=self.device)
            else:
                if self._config.save_vram:
                    return torch.zeros(
                        (size, *dim), device="cpu"
                    )  # Move to CPU to avoid GPU OOM, but will be slower
                else:
                    return torch.zeros((size, *dim), device=self.device)

        def make_numpy_buffer(size: int):
            return np.empty((size,), dtype=object)

        # Always create sample ID buffers
        self._sample_ids_buffer = {
            "train": make_numpy_buffer(self.n_train),
            "valid": make_numpy_buffer(self.n_valid) if self.n_valid else None,
            "test": make_numpy_buffer(self.n_test) if self.n_test else None,
        }

        splits = ["test"] if self._config.save_memory else ["train", "valid", "test"]

        def make_split_buffers(dim):
            return {
                split: (
                    make_tensor_buffer(getattr(self, f"n_{split}"), dim)
                    if getattr(self, f"n_{split}", 0) and (split in splits)
                    else None
                )
                for split in ["train", "valid", "test"]
            }

        self._latentspace_buffer = make_split_buffers(self.latent_dim)
        self._reconstruction_buffer = make_split_buffers(self.n_features)
        self._mu_buffer = make_split_buffers(self.latent_dim)
        self._sigma_buffer = make_split_buffers(self.latent_dim)

    def _ontix_hook(self):
        pass

    def train(self, epochs_overwrite=None) -> Result:
        """Orchestrates training over multiple epochs, including training and validation phases.

        Args:
            epochs_overwrite: If provided, overrides the number of epochs specified in the config.
                This is only there so we can use the train method for pretraining.Any
        Returns:
            Result object containing training results and dynamics like latent spaces and reconstructions.
        """
        epochs = self._config.epochs
        if epochs_overwrite:
            epochs = epochs_overwrite
        with self._fabric.autocast():
            for epoch in range(epochs):
                self._init_buffers()
                should_checkpoint: bool = self._should_checkpoint(epoch)
                self._model.train()

                epoch_loss, epoch_sub_losses = self._train_epoch(
                    should_checkpoint=should_checkpoint, epoch=epoch
                )
                self._log_losses(epoch, "train", epoch_loss, epoch_sub_losses)

                if self._validset:
                    valid_loss, valid_sub_losses = self._validate_epoch(
                        should_checkpoint=should_checkpoint, epoch=epoch
                    )
                    self._log_losses(epoch, "valid", valid_loss, valid_sub_losses)

                if should_checkpoint:
                    self._store_checkpoint(epoch)

        self._result.model = next(self._model.children())
        return self._result

    def maskix_hook(self, X: torch.Tensor):
        """Only active when override from MaskixTrainer is used"""
        return X

    def _train_epoch(
        self, should_checkpoint: bool, epoch: int
    ) -> Tuple[float, Dict[str, float]]:
        """Handles loss computation, backwards pass and checkpointing for one epoch.

        Args:
            should_checkpoint: Whether to checkpoint this epoch.
            epoch: The current epoch number.
        Returns:
            total_loss: The total loss for the epoch.
            sub_losses: A dictionary of sub-losses accumulated over the epoch.
        """
        total_loss = 0.0
        sub_losses: Dict[str, float] = defaultdict(float)
        current_batch = 0
        for indices, features, sample_ids in self._trainloader:
            current_batch += 1
            self._optimizer.zero_grad()
            X: torch.Tensor = self.maskix_hook(features)
            model_outputs = self._model(X)
            loss, batch_sub_losses = self._loss_fn(
                model_output=model_outputs,
                targets=features,
                corrupted_input=X,  # only relevant for MaskixLoss
                epoch=epoch,
                n_samples=len(
                    self._trainloader.dataset
                ),  # Pass n_samples for disentangled loss calculations
            )
            self._fabric.backward(loss)
            self._ontix_hook()
            self._optimizer.step()

            total_loss += loss.item()
            for k, v in batch_sub_losses.items():
                if "_factor" not in k:  # Skip factor losses
                    sub_losses[k] += v.item()
                else:
                    sub_losses[k] = v.item()

            if should_checkpoint:
                self._capture_dynamics(model_outputs, "train", indices, sample_ids)

        for k, v in sub_losses.items():
            if "_factor" not in k:
                sub_losses[k] = v / len(
                    self._trainloader.dataset
                )  # Average over all samples
        total_loss = total_loss / len(
            self._trainloader.dataset
        )  # Average over all samples
        return total_loss, sub_losses

    def _validate_epoch(
        self, should_checkpoint: bool, epoch: int
    ) -> Tuple[float, Dict[str, float]]:
        """Handles validation for one epoch.

        Args:
            should_checkpoint: Whether to checkpoint this epoch.
            epoch: The current epoch number.
        Returns:
            total_loss: The total loss for the epoch.
            sub_losses: A dictionary of sub-losses accumulated over the epoch.
        """
        total_loss = 0.0
        sub_losses: Dict[str, float] = defaultdict(float)
        self._model.eval()

        with torch.no_grad():
            for indices, features, sample_ids in self._validloader:
                X: torch.Tensor = self.maskix_hook(features)
                model_outputs = self._model(X)
                loss, batch_sub_losses = self._loss_fn(
                    model_output=model_outputs,
                    targets=features,
                    corrupted_input=X,
                    epoch=epoch,
                    n_samples=len(
                        self._validloader.dataset
                    ),  # Pass n_samples for disentangled loss calculations
                )
                total_loss += loss.item()
                for k, v in batch_sub_losses.items():
                    if "_factor" not in k:  # Skip factor losses
                        sub_losses[k] += v.item()
                    else:
                        sub_losses[k] = v.item()
                if should_checkpoint:
                    self._capture_dynamics(model_outputs, "valid", indices, sample_ids)

        for k, v in sub_losses.items():
            if "_factor" not in k:
                sub_losses[k] = v / len(
                    self._validloader.dataset
                )  # Average over all samples
        total_loss = total_loss / len(
            self._validloader.dataset
        )  # Average over all samples
        return total_loss, sub_losses

    def _log_losses(
        self, epoch: int, split: str, total_loss: float, sub_losses: Dict[str, float]
    ) -> None:
        """Logs the total and sub-losses for an epoch and stores them in the Result object.

        Args:
            epoch: The current epoch number.
            split: The data split ("train" or "valid").
            total_loss: The total loss for the epoch.
            sub_losses: A dictionary of sub-losses for the epoch.
        """
        # dataset_len = len(
        #     self._trainloader.dataset if split == "train" else self._validloader.dataset
        # )
        self._result.losses.add(epoch=epoch, split=split, data=total_loss)
        self._result.sub_losses.add(
            epoch=epoch,
            split=split,
            data={k: v if "_factor" not in k else v for k, v in sub_losses.items()},
        )

        self._fabric.print(
            f"Epoch {epoch + 1} - {split.capitalize()} Loss: {total_loss:.4f}"
        )
        self._fabric.print(
            f"Sub-losses: {', '.join([f'{k}: {v:.4f}' for k, v in sub_losses.items()])}"
        )

    def _store_checkpoint(self, epoch: int) -> None:
        """Stores model checkpoints and training dynamics to result object.

        Args:
            epoch: The current epoch number.
        """
        self._result.model_checkpoints.add(epoch=epoch, data=self._model.state_dict())
        if self._config.save_memory:
            return
        self._dynamics_to_result(epoch, "train")
        if self._validset:
            self._dynamics_to_result(epoch, "valid")

    def _capture_dynamics(
        self,
        model_output: Union[ModelOutput, Any],
        split: str,
        indices: torch.Tensor,
        sample_ids: Any,
        **kwargs,
    ) -> None:
        """Writes model dynamics (latent space, reconstructions, etc.) to buffers.

        Args:
            model_output: The output from the model containing dynamics information.
            split: The data split ("train", "valid", or "test").
            indices: The indices of the samples in the current batch.
            sample_ids: The sample IDs corresponding to the current batch.
            **kwargs: Additional arguments (not used here).

        """
        indices_np = (
            indices.cpu().numpy()
            if isinstance(indices, torch.Tensor)
            else np.array(indices)
        )
        if isinstance(sample_ids, torch.Tensor):
            sample_ids = sample_ids.cpu().numpy()
        self._sample_ids_buffer[split][indices_np] = np.array(sample_ids)
        if self._config.save_memory and split != "test":
            return
        indices_np = (
            indices.cpu().numpy()
            if isinstance(indices, torch.Tensor)
            else np.array(indices)
        )

        self._sample_ids_buffer[split][indices_np] = np.array(sample_ids)

        if self._config.save_vram:  # Move to CPU to avoid GPU OOM, but will be slower
            self._latentspace_buffer[split][
                indices_np
            ] = model_output.latentspace.cpu().detach()
            self._reconstruction_buffer[split][
                indices_np
            ] = model_output.reconstruction.cpu().detach()
            if model_output.latent_logvar is not None:
                self._sigma_buffer[split][
                    indices_np
                ] = model_output.latent_logvar.cpu().detach()
            if model_output.latent_mean is not None:
                self._mu_buffer[split][
                    indices_np
                ] = model_output.latent_mean.cpu().detach()
        else:
            self._latentspace_buffer[split][
                indices_np
            ] = model_output.latentspace.detach()
            self._reconstruction_buffer[split][
                indices_np
            ] = model_output.reconstruction.detach()
            if model_output.latent_logvar is not None:
                self._sigma_buffer[split][
                    indices_np
                ] = model_output.latent_logvar.detach()
            if model_output.latent_mean is not None:
                self._mu_buffer[split][indices_np] = model_output.latent_mean.detach()

    def _dynamics_to_result(self, epoch: int, split: str) -> None:
        """Transfers buffered dynamics to the Result object.

        Args:
            epoch: The current epoch number.
            split: The data split ("train", "valid", or "test").
        """

        def maybe_add(buffer, target):
            if buffer[split] is not None and buffer[split].sum() != 0:
                target.add(
                    epoch=epoch, split=split, data=buffer[split].cpu().detach().numpy()
                )

        self._result.latentspaces.add(
            epoch=epoch,
            split=split,
            data=self._latentspace_buffer[split].cpu().detach().numpy(),
        )
        self._result.reconstructions.add(
            epoch=epoch,
            split=split,
            data=self._reconstruction_buffer[split].cpu().detach().numpy(),
        )
        self._result.sample_ids.add(
            epoch=epoch, split=split, data=self._sample_ids_buffer[split]
        )
        maybe_add(self._mu_buffer, self._result.mus)
        maybe_add(self._sigma_buffer, self._result.sigmas)

    def predict(
        self, data: BaseDataset, model: Optional[torch.nn.Module] = None, **kwargs
    ) -> Result:
        """Decided to add predict method to the trainer class.

        This violates SRP, but the trainer class has a lot of attributes and methods
        that are needed for prediction. So this way we don't need to write so much duplicate code

        Args:
            data: BaseDataset unseen data to run inference on
            model: torch.nn.Module model to run inference with
            **kwargs: Additional arguments (not used here).

        Returns:
            self._result: Result object containing the inference results

        """
        if model is None:
            model: torch.nn.Module = self._model
            import warnings

            warnings.warn(
                "No model provided for prediction, using the trained model from the trainer."
            )
        model.eval()
        inference_loader = DataLoader(
            data,
            batch_size=self._config.batch_size,
            shuffle=False,
        )
        self.n_test = len(data)
        self._init_buffers(input_data=data)
        inference_loader = self._fabric.setup_dataloaders(inference_loader)  # type: ignore
        with self._fabric.autocast(), torch.inference_mode():
            processed_samples = 0
            for idx, data, sample_ids in inference_loader:
                model_output = model(data)
                processed_samples += len(data)
                print(
                    f"Processed {processed_samples} / {self.n_test} samples", end="\r"
                )
                self._capture_dynamics(
                    model_output=model_output,
                    split="test",
                    sample_ids=sample_ids,
                    indices=idx,
                )
        self._dynamics_to_result(epoch=-1, split="test")

        return self._result

    def decode(self, x: torch.Tensor):
        """Decodes the input tensor x using the trained model.

        Args:
            x: Input tensor to be decoded.
        Returns:
            Decoded tensor.
        """
        with self._fabric.autocast(), torch.no_grad():
            x = self._fabric.to_device(obj=x)
            if not isinstance(x, torch.Tensor):
                raise TypeError(
                    f"Expected input to be a torch.Tensor, got {type(x)} instead."
                )
            return self._model.decode(x=x)

    def get_model(self) -> torch.nn.Module:
        """Getter method for the trained model.

        Returns:

            The trained model as a torch.nn.Module.
        """
        return self._model

    def purge(self) -> None:
        """Cleans up any resources used during training, such as cached data or large attributes."""

        attrs_to_delete = [
            "_trainloader",
            "_validloader",
            "_model",
            "_optimizer",
            "_latentspace_buffer",
            "_reconstruction_buffer",
            "_mu_buffer",
            "_sigma_buffer",
            "_sample_ids_buffer",
            "_trainset",
            "_loss_fn",
            "_validset",
        ]

        for attr in attrs_to_delete:
            if hasattr(self, attr):
                value = getattr(self, attr)
                # Optional: ensure non-None before deletion
                if value is not None:
                    delattr(self, attr)

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        gc.collect()

__init__(trainset, validset, result, config, model_type, loss_type, ontologies=None, **kwargs)

Initializes the GeneralTrainer.

Parameters:

Name	Type	Description	Default
trainset	Optional[BaseDataset]	The dataset used for training.	required
validset	Optional[BaseDataset]	The dataset used for validation, if provided.	required
result	Result	An object to store and manage training results.	required
config	DefaultConfig	Configuration object containing training hyperparameters and settings.	required
model_type	Type[BaseAutoencoder]	The autoencoder model class to be trained.	required
loss_type	Type[BaseLoss]	The loss function class to be used for training.	required
ontologies	Optional[Union[Tuple, List]]	Ontology information, if provided for Ontix	None

Source code in src/autoencodix/trainers/_general_trainer.py

def __init__(
    self,
    trainset: Optional[BaseDataset],
    validset: Optional[BaseDataset],
    result: Result,
    config: DefaultConfig,
    model_type: Type[BaseAutoencoder],
    loss_type: Type[BaseLoss],
    ontologies: Optional[Union[Tuple, List]] = None,
    **kwargs,
):
    """Initializes the GeneralTrainer.

    Args:
            trainset: The dataset used for training.
            validset: The dataset used for validation, if provided.
            result: An object to store and manage training results.
            config: Configuration object containing training hyperparameters and settings.
            model_type: The autoencoder model class to be trained.
            loss_type: The loss function class to be used for training.
            ontologies: Ontology information, if provided for Ontix
    """
    self._n_cpus = os.cpu_count()
    if self._n_cpus is None:
        self._n_cpus = 0

    super().__init__(
        trainset, validset, result, config, model_type, loss_type, ontologies
    )

    self.latent_dim = config.latent_dim
    if ontologies is not None:
        if not hasattr(self._model, "latent_dim"):
            raise ValueError(
                "Model must have a 'latent_dim' attribute when ontologies are provided."
            )
        self.latent_dim = self._model.latent_dim

    # we will set this later, in the predict method
    self.n_test: Optional[int] = None
    self.n_train = len(trainset) if trainset else 0
    self.n_valid = len(validset) if validset else 0
    self.n_features = trainset.get_input_dim() if trainset else 0
    self.device = next(self._model.parameters()).device
    self._init_buffers()

decode(x)

Decodes the input tensor x using the trained model.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor to be decoded.	required

Returns: Decoded tensor.

Source code in src/autoencodix/trainers/_general_trainer.py

def decode(self, x: torch.Tensor):
    """Decodes the input tensor x using the trained model.

    Args:
        x: Input tensor to be decoded.
    Returns:
        Decoded tensor.
    """
    with self._fabric.autocast(), torch.no_grad():
        x = self._fabric.to_device(obj=x)
        if not isinstance(x, torch.Tensor):
            raise TypeError(
                f"Expected input to be a torch.Tensor, got {type(x)} instead."
            )
        return self._model.decode(x=x)

get_model()

Getter method for the trained model.

Returns:

The trained model as a torch.nn.Module.

Source code in src/autoencodix/trainers/_general_trainer.py

def get_model(self) -> torch.nn.Module:
    """Getter method for the trained model.

    Returns:

        The trained model as a torch.nn.Module.
    """
    return self._model

maskix_hook(X)

Only active when override from MaskixTrainer is used

Source code in src/autoencodix/trainers/_general_trainer.py

def maskix_hook(self, X: torch.Tensor):
    """Only active when override from MaskixTrainer is used"""
    return X

predict(data, model=None, **kwargs)

Decided to add predict method to the trainer class.

This violates SRP, but the trainer class has a lot of attributes and methods that are needed for prediction. So this way we don't need to write so much duplicate code

Parameters:

Name	Type	Description	Default
data	BaseDataset	BaseDataset unseen data to run inference on	required
model	Optional[Module]	torch.nn.Module model to run inference with	None
**kwargs		Additional arguments (not used here).	{}

Returns:

Type	Description
Result	self._result: Result object containing the inference results

Source code in src/autoencodix/trainers/_general_trainer.py

def predict(
    self, data: BaseDataset, model: Optional[torch.nn.Module] = None, **kwargs
) -> Result:
    """Decided to add predict method to the trainer class.

    This violates SRP, but the trainer class has a lot of attributes and methods
    that are needed for prediction. So this way we don't need to write so much duplicate code

    Args:
        data: BaseDataset unseen data to run inference on
        model: torch.nn.Module model to run inference with
        **kwargs: Additional arguments (not used here).

    Returns:
        self._result: Result object containing the inference results

    """
    if model is None:
        model: torch.nn.Module = self._model
        import warnings

        warnings.warn(
            "No model provided for prediction, using the trained model from the trainer."
        )
    model.eval()
    inference_loader = DataLoader(
        data,
        batch_size=self._config.batch_size,
        shuffle=False,
    )
    self.n_test = len(data)
    self._init_buffers(input_data=data)
    inference_loader = self._fabric.setup_dataloaders(inference_loader)  # type: ignore
    with self._fabric.autocast(), torch.inference_mode():
        processed_samples = 0
        for idx, data, sample_ids in inference_loader:
            model_output = model(data)
            processed_samples += len(data)
            print(
                f"Processed {processed_samples} / {self.n_test} samples", end="\r"
            )
            self._capture_dynamics(
                model_output=model_output,
                split="test",
                sample_ids=sample_ids,
                indices=idx,
            )
    self._dynamics_to_result(epoch=-1, split="test")

    return self._result

purge()

Cleans up any resources used during training, such as cached data or large attributes.

Source code in src/autoencodix/trainers/_general_trainer.py

def purge(self) -> None:
    """Cleans up any resources used during training, such as cached data or large attributes."""

    attrs_to_delete = [
        "_trainloader",
        "_validloader",
        "_model",
        "_optimizer",
        "_latentspace_buffer",
        "_reconstruction_buffer",
        "_mu_buffer",
        "_sigma_buffer",
        "_sample_ids_buffer",
        "_trainset",
        "_loss_fn",
        "_validset",
    ]

    for attr in attrs_to_delete:
        if hasattr(self, attr):
            value = getattr(self, attr)
            # Optional: ensure non-None before deletion
            if value is not None:
                delattr(self, attr)

    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    gc.collect()

train(epochs_overwrite=None)

Orchestrates training over multiple epochs, including training and validation phases.

Parameters:

Name	Type	Description	Default
epochs_overwrite		If provided, overrides the number of epochs specified in the config. This is only there so we can use the train method for pretraining.Any	None

Returns: Result object containing training results and dynamics like latent spaces and reconstructions.

Source code in src/autoencodix/trainers/_general_trainer.py

def train(self, epochs_overwrite=None) -> Result:
    """Orchestrates training over multiple epochs, including training and validation phases.

    Args:
        epochs_overwrite: If provided, overrides the number of epochs specified in the config.
            This is only there so we can use the train method for pretraining.Any
    Returns:
        Result object containing training results and dynamics like latent spaces and reconstructions.
    """
    epochs = self._config.epochs
    if epochs_overwrite:
        epochs = epochs_overwrite
    with self._fabric.autocast():
        for epoch in range(epochs):
            self._init_buffers()
            should_checkpoint: bool = self._should_checkpoint(epoch)
            self._model.train()

            epoch_loss, epoch_sub_losses = self._train_epoch(
                should_checkpoint=should_checkpoint, epoch=epoch
            )
            self._log_losses(epoch, "train", epoch_loss, epoch_sub_losses)

            if self._validset:
                valid_loss, valid_sub_losses = self._validate_epoch(
                    should_checkpoint=should_checkpoint, epoch=epoch
                )
                self._log_losses(epoch, "valid", valid_loss, valid_sub_losses)

            if should_checkpoint:
                self._store_checkpoint(epoch)

    self._result.model = next(self._model.children())
    return self._result

StackixOrchestrator

StackixOrchestrator coordinates the training of multi-modality VAE stacking.

This orchestrator manages both parallel and sequential training of modality-specific autoencoders, followed by creating a concatenated latent space for the final stacked model training. It leverages Lightning Fabric's distribution strategies for efficient training.

Attributes: _workdir: Directory for saving intermediate models and results modality_models: Dictionary of trained models for each modality modality_results: Dictionary of training results for each modality _modality_latent_dims: Dictionary of latent dimensions for each modality concatenated_latent_spaces: Dictionary of concatenated latent spaces by split _dataset_type: Class to use for creating datasets _fabric: Lightning Fabric instance for distributed operations _trainer_class: Class to use for training (dependency injection) trainset: Training dataset containing multiple modalities validset: Validation dataset containing multiple modalities testset: Test dataset containing multiple modalities loss_type: Type of loss function to use for training model_type: Type of autoencoder model to use for each modality config: Configuration parameters for training and model architecture stacked_model: The final stacked autoencoder model (initialized later) stacked_trainer: Trainer for the stacked model (initialized later) concat_idx: Dictionary tracking the start and end indices of each modality in the concatenated latent space dropped_indices_map: Dictionary tracking dropped sample indices for each modality during alignment reconstruction_shapes: Dictionary storing original shapes of latent spaces for reconstruction common_sample_ids: Common sample IDs across all modalities for alignment

Source code in src/autoencodix/trainers/_stackix_orchestrator.py

class StackixOrchestrator:
    """StackixOrchestrator coordinates the training of multi-modality VAE stacking.

    This orchestrator manages both parallel and sequential training of modality-specific
    autoencoders, followed by creating a concatenated latent space for the final
    stacked model training. It leverages Lightning Fabric's distribution strategies
    for efficient training.

    Attributes:
    _workdir: Directory for saving intermediate models and results
    modality_models: Dictionary of trained models for each modality
    modality_results: Dictionary of training results for each modality
    _modality_latent_dims: Dictionary of latent dimensions for each modality
    concatenated_latent_spaces: Dictionary of concatenated latent spaces by split
    _dataset_type: Class to use for creating datasets
    _fabric: Lightning Fabric instance for distributed operations
    _trainer_class: Class to use for training (dependency injection)
    trainset: Training dataset containing multiple modalities
    validset: Validation dataset containing multiple modalities
    testset: Test dataset containing multiple modalities
    loss_type: Type of loss function to use for training
    model_type: Type of autoencoder model to use for each modality
    config: Configuration parameters for training and model architecture
    stacked_model: The final stacked autoencoder model (initialized later)
    stacked_trainer: Trainer for the stacked model (initialized later)
    concat_idx: Dictionary tracking the start and end indices of each modality in the concatenated latent space
    dropped_indices_map: Dictionary tracking dropped sample indices for each modality during alignment
    reconstruction_shapes: Dictionary storing original shapes of latent spaces for reconstruction
    common_sample_ids: Common sample IDs across all modalities for alignment

    """

    def __init__(
        self,
        trainset: Optional[MultiModalDataset],
        validset: Optional[MultiModalDataset],
        config: DefaultConfig,
        model_type: Type[BaseAutoencoder],
        loss_type: Type[BaseLoss],
        testset: Optional[MultiModalDataset] = None,
        trainer_type: Type[GeneralTrainer] = GeneralTrainer,
        dataset_type: Type[BaseDataset] = NumericDataset,
        workdir: str = "./stackix_work",
    ):
        """
        Initialize the StackixOrchestrator with datasets and configuration.

        Args:
            trainset: Training dataset containing multiple modalities
            validset: Validation dataset containing multiple modalities
            config: Configuration parameters for training and model architecture
            model_type: Type of autoencoder model to use for each modality
            loss_type: Type of loss function to use for training
            testset: Dataset with test split
            trainer_type: Type to use for training (default is GeneralTrainer)
            dataset_type: Type to use for creating datasets (default is NumericDataset)
            workdir: Directory to save intermediate models and results (default is "./stackix_work")
        """
        self._trainset = trainset
        self._validset = validset
        self._testset = testset
        self._config = config
        self._model_type = model_type
        self._loss_type = loss_type
        self._workdir = workdir
        self._trainer_class = trainer_type
        self._dataset_type = dataset_type

        # Initialize fabric for distributed operations
        strategy = "auto"
        if hasattr(config, "gpu_strategy"):
            strategy = config.gpu_strategy

        self._fabric = lightning_fabric.Fabric(
            accelerator=(
                str(config.device) if hasattr(config, "device") else None
            ),  # ty: ignore[invalid-argument-type]
            devices=config.n_gpus if hasattr(config, "n_gpus") else 1,
            precision=(
                config.float_precision if hasattr(config, "float_precision") else "32"
            ),
            strategy=strategy,
        )

        # Initialize storage for models and results
        self.modality_trainers: Dict[str, BaseAutoencoder] = {}
        self.modality_results: Dict[str, Result] = {}
        self._modality_latent_dims: Dict[str, int] = {}
        self.concatenated_latent_spaces: Dict[str, torch.Tensor] = {}

        # Stacked model and trainer will be created during the process
        self.stacked_model: Optional[BaseAutoencoder] = None
        self.stacked_trainer: Optional[GeneralTrainer] = None
        self.concat_idx: Dict[str, Optional[Tuple[int, int]]] = {}
        self.dropped_indices_map: Dict[str, np.ndarray] = {}
        self.reconstruction_shapes: Dict[str, torch.Size] = {}
        self.common_sample_ids: Optional[pd.Index] = None

        # Create the directory if it doesn't exist
        os.makedirs(name=self._workdir, exist_ok=True)

    def set_testset(self, testset: MultiModalDataset) -> None:
        """Set the test dataset for the orchestrator.

        Args:
            testset: Test dataset containing multiple modalities
        """
        self._testset = testset

    def _train_modality(
        self,
        modality: str,
        modality_dataset: BaseDataset,
        valid_modality_dataset: Optional[BaseDataset] = None,
    ) -> None:
        """Trains a single modality and returns the trained model and result.

        This function is designed to be executed within a Lightning Fabric context.
        It trains an individual modality model and saves the results to disk.

        Args:
            modality: Modality name/identifier
            modality_dataset: Training dataset for this modality
            valid_modality_dataset: Validation dataset for this modality (default is None)

        Returns:
            Trained model and result object
        """
        print(f"Training modality: {modality}")
        result = Result()
        trainer = self._trainer_class(
            trainset=modality_dataset,
            validset=valid_modality_dataset,
            result=result,
            config=self._config,
            model_type=self._model_type,
            loss_type=self._loss_type,
        )

        result = trainer.train()
        self.modality_results[modality] = result
        self.modality_trainers[modality] = trainer

    def _train_distributed(self, keys: List[str]) -> None:
        """Trains modality models in a distributed fashion using Lightning Fabric.

        Uses Lightning Fabric's built-in capabilities for distributing work across devices.
        Each process trains a subset of modalities, then loads results from other processes.

        Args:
            keys: List of modality keys to train
        """
        # For parallel training with multiple devices
        strategy = (
            self._config.gpu_strategy
            if hasattr(self._config, "gpu_strategy")
            else "auto"
        )
        n_gpus = self._config.n_gpus if hasattr(self._config, "n_gpus") else 1
        device: str = self._config.device

        if strategy == "auto" and device in ["cuda", "gpu"]:
            strategy = "ddp" if n_gpus > 1 else "dp"

        # Create a fabric instance with appropriate parallelism strategy
        fabric = lightning_fabric.Fabric(
            accelerator=device,
            devices=min(n_gpus, len(keys)),  # No more devices than modalities
            precision=(
                self._config.float_precision
                if hasattr(self._config, "float_precision")
                else "32"
            ),
            strategy=strategy,
        )

        # Launch distributed training
        fabric.launch()

        # Device ID within the process group
        local_rank = fabric.local_rank if hasattr(fabric, "local_rank") else 0
        world_size = fabric.world_size if hasattr(fabric, "world_size") else 1

        # Distribute modalities across ranks
        my_modalities = [k for i, k in enumerate(keys) if i % world_size == local_rank]

        # Train assigned modalities
        for modality in my_modalities:
            train_dataset = self._trainset.datasets[modality]
            valid_dataset = (
                self._validset.datasets.get(modality)
                if self._validset and hasattr(self._validset, "datasets")
                else None
            )

            self._train_modality(
                modality=modality,
                modality_dataset=train_dataset,
                valid_modality_dataset=valid_dataset,
            )
            self._modality_latent_dims[modality] = train_dataset.get_input_dim()

        # Synchronize across processes to ensure all modalities are trained
        if world_size > 1:
            fabric.barrier()

        # Load models and results from other processes
        for modality in keys:
            if modality in self.modality_trainers.keys():
                continue
            # Load model state dict
            model_path = os.path.join(self._workdir, f"{modality}_model.ckpt")

            # Get input dimension for this modality
            input_dim = self._trainset.datasets[modality].get_input_dim()

            # Initialize a new model
            model = self._model_type(config=self._config, input_dim=input_dim)
            model.load_state_dict(
                state_dict=torch.load(f=model_path, map_location="cpu")
            )

            # Load result
            with open(
                file=os.path.join(self._workdir, f"{modality}_result.pkl"),
                mode="rb",
            ) as f:
                result = pickle.load(file=f)

            self.modality_trainers[modality] = model
            self.modality_results[modality] = result
            self._modality_latent_dims[modality] = input_dim

    def _train_sequential(self, keys: List[str]) -> None:
        """Trains modality models sequentially on a single device.

        Used when distributed training is not available or not necessary.
        Processes each modality one after another.

        Args:
            keys: List of modality keys to train
        """
        for modality in keys:
            print(f"Training modality: {modality}")
            train_dataset = self._trainset.datasets[modality]
            valid_dataset = (
                self._validset.datasets.get(modality) if self._validset else None
            )

            self._train_modality(
                modality=modality,
                modality_dataset=train_dataset,
                valid_modality_dataset=valid_dataset,
            )

            self._modality_latent_dims[modality] = train_dataset.get_input_dim()

    def _extract_latent_spaces(
        self, result_dict: Dict[str, Any], split: str = "train"
    ) -> torch.Tensor:
        """Extracts, aligns, and concatenates latent spaces, populating instance attributes for later reconstruction.

        Args:
            result_dict: Dictionary of Result objects from modality training
            split: Data split to extract latent spaces from ("train", "valid", "test"), default is "train"
        """
        # --- Step 1: Extract Latent Spaces and Sample IDs ---
        all_latents = {}
        all_sample_ids = {}
        for name, result in result_dict.items():
            try:
                latent_space = result.latentspaces.get(split=split, epoch=-1)
                sample_ids = result.sample_ids.get(split=split, epoch=-1)

                if latent_space.shape[0] != len(sample_ids):
                    raise ValueError(
                        f"Mismatch between latent space rows ({latent_space.shape[0]}) and sample IDs ({len(sample_ids)}) for modality '{name}'."
                    )

                all_latents[name] = latent_space
                all_sample_ids[name] = sample_ids
                self.reconstruction_shapes[name] = latent_space.shape
            except Exception as e:
                raise ValueError(
                    f"Failed to extract data for modality {name} on split {split}: {e}"
                )

        if not all_latents:
            raise ValueError("No latent spaces were extracted.")

        # --- Step 2: Find Common Sample IDs ---
        initial_ids = set(next(iter(all_sample_ids.values())))
        common_ids_set = initial_ids.intersection(
            *[set(ids) for ids in all_sample_ids.values()]
        )

        if not common_ids_set:
            raise ValueError(
                "No common samples found across all modalities for concatenation."
            )
        # common_ids = [cid for cid in list(common_ids_set) if cid]
        self.common_sample_ids = pd.Index(sorted(list(common_ids_set)))

        # self.common_sample_ids = pd.Index(sorted(list(common_ids_set)))
        print(
            f"Found {len(self.common_sample_ids)} common samples for the stacked autoencoder."
        )

        # --- Step 3 & 4: Align Latent Spaces and Track Dropped Indices ---
        aligned_latents_list = []
        start_idx = 0
        for name, latent_space in all_latents.items():
            original_ids = pd.Index(all_sample_ids[name])
            latent_df = pd.DataFrame(latent_space, index=original_ids)

            aligned_df = latent_df.loc[self.common_sample_ids]
            aligned_latents_list.append(torch.from_numpy(aligned_df.values))

            end_idx = start_idx + aligned_df.shape[1]
            self.concat_idx[name] = (start_idx, end_idx)
            start_idx = end_idx

            dropped_mask = ~original_ids.isin(self.common_sample_ids)
            self.dropped_indices_map[name] = np.where(dropped_mask)[0]

        # --- Step 5: Concatenate ---
        stackix_input = torch.cat(aligned_latents_list, dim=1)
        return stackix_input

    def train_modalities(self) -> Tuple[Dict[str, BaseAutoencoder], Dict[str, Result]]:
        """Trains all modality-specific models.

        This is the first phase of Stackix training where each modality is
        trained independently before their latent spaces are combined.

        Returns:
            Dictionary of trained models for each modality and Dictionary of training results
            for each modality.

        Raises:
            ValueError: If trainset is not a MultiModalDataset or has no modalities
        """
        # if not isinstance(self._trainset, MulDataset):
        #     raise ValueError("Trainset must be a MultiModalDataset for Stackix training")

        keys = self._trainset.datasets.keys()
        if not keys:
            raise ValueError("No modalities found in trainset")

        n_modalities = len(keys)

        # Determine if we should use distributed training
        n_gpus = self._config.n_gpus if hasattr(self._config, "n_gpus") else 0
        use_distributed = bool(n_gpus and n_gpus > 1 and n_modalities > 1)

        if use_distributed:
            self._train_distributed(keys=keys)
        else:
            self._train_sequential(keys=keys)

        return self.modality_trainers, self.modality_results

    def prepare_latent_datasets(self, split: str) -> NumericDataset:
        """Prepares datasets with concatenated latent spaces for stacked model training.

        This is the second phase of Stackix training where latent spaces from
        all modalities are extracted and concatenated.

        Returns:
            Training and validation datasets with concatenated latent spaces

        Raises:
            ValueError: If no modality models have been trained or no latent spaces could be extracted
        """
        if not self.modality_trainers:
            raise ValueError(
                "No modality models have been trained. Call train_modalities() first."
            )

        # Extract and concatenate latent spaces
        if split == "test":
            if self._testset is None:
                raise ValueError(
                    "No test dataset available. Please provide a test dataset for prediction."
                )
            latent = self._extract_latent_spaces(
                result_dict=self.predict_modalities(data=self._testset), split=split
            )
            ds = NumericDataset(
                data=latent,
                config=self._config,
                sample_ids=self.common_sample_ids,
                feature_ids=[
                    f"{k}_latent_{i}"
                    for k, (start, end) in self.concat_idx.items()
                    for i in range(start, end)
                ],
            )
            return ds

        latent = self._extract_latent_spaces(
            result_dict=self.modality_results, split=split
        )
        # Create datasets for the concatenated latent spaces
        feature_ids = [
            f"{k}_latent_{i}"
            for k, (start, end) in self.concat_idx.items()
            for i in range(start, end)
        ]
        ds = NumericDataset(
            data=latent,
            config=self._config,
            sample_ids=self.common_sample_ids,
            feature_ids=feature_ids,
        )

        return ds

    def predict_modalities(self, data: MultiModalDataset) -> Dict[str, torch.Tensor]:
        """Predicts using the trained models for each modality.

        Args:
            data: Input data for prediction, uses test data if not provided

        Returns:
            Dictionary of reconstructed tensors by modality

        Raises:
            ValueError: If model has not been trained yet or no data is available
        """
        if not self.modality_trainers:
            raise ValueError(
                "No modality models have been trained. Call train_modalities() first."
            )

        predictions = {}
        for modality, trainer in self.modality_trainers.items():
            predictions[modality] = trainer.predict(
                data=data.datasets[modality],
                model=trainer._model,
            )
        return predictions

    def reconstruct_from_stack(
        self, reconstructed_stack: torch.Tensor
    ) -> Dict[str, torch.Tensor]:
        """Reconstructs the full data for each modality from the stacked latent reconstruction.

        Args:
            reconstructed_stack: Tensor with the reconstructed concatenated latent space
        Returns:
            Dictionary of reconstructed tensors by modality
        """
        modality_reconstructions: Dict[str, torch.Tensor] = {}

        for trainer in self.modality_trainers.values():
            trainer._model.eval()

        for name, (start_idx, end_idx) in self.concat_idx.items():
            # =================================================================
            # STEP 1: DE-CONCATENATE THE ALIGNED PART
            # This is the small, dense reconstruction matching the common samples.
            # =================================================================
            aligned_latent_recon = reconstructed_stack[:, start_idx:end_idx]

            # =================================================================
            # STEP 2: RE-ASSEMBLE THE FULL-SIZE LATENT SPACE
            # This is the logic from the old `reconstruct_full_latent` function,
            # now integrated directly here.
            # =================================================================
            original_shape = self.reconstruction_shapes[name]
            dropped_indices = self.dropped_indices_map[name]
            n_original_samples = original_shape[0]

            # Determine the original integer positions of the samples that were KEPT.
            kept_indices = np.delete(np.arange(n_original_samples), dropped_indices)

            # Create a full-size placeholder tensor matching the original latent space.
            # We use zeros, as they are a neutral input for most decoders.
            full_latent_recon = torch.zeros(
                size=original_shape,
                dtype=aligned_latent_recon.dtype,
                device=self._fabric.device,
            )

            # Place the aligned reconstruction data into the correct rows of the full tensor.
            full_latent_recon[kept_indices, :] = self._fabric.to_device(
                aligned_latent_recon
            )

            # =================================================================
            # STEP 3: DECODE THE FULL-SIZE LATENT TENSOR
            # The decoder now receives a tensor with the correct number of samples,
            # including the rows of zeros for the dropped samples.
            # =================================================================
            with torch.no_grad():
                model = self.modality_trainers[name].get_model()
                model = self._fabric.to_device(model)

                final_recon = model.decode(full_latent_recon)
                modality_reconstructions[name] = final_recon.cpu()

        return modality_reconstructions

__init__(trainset, validset, config, model_type, loss_type, testset=None, trainer_type=GeneralTrainer, dataset_type=NumericDataset, workdir='./stackix_work')

Initialize the StackixOrchestrator with datasets and configuration.

Parameters:

Name	Type	Description	Default
trainset	Optional[MultiModalDataset]	Training dataset containing multiple modalities	required
validset	Optional[MultiModalDataset]	Validation dataset containing multiple modalities	required
config	DefaultConfig	Configuration parameters for training and model architecture	required
model_type	Type[BaseAutoencoder]	Type of autoencoder model to use for each modality	required
loss_type	Type[BaseLoss]	Type of loss function to use for training	required
testset	Optional[MultiModalDataset]	Dataset with test split	None
trainer_type	Type[GeneralTrainer]	Type to use for training (default is GeneralTrainer)	GeneralTrainer
dataset_type	Type[BaseDataset]	Type to use for creating datasets (default is NumericDataset)	NumericDataset
workdir	str	Directory to save intermediate models and results (default is "./stackix_work")	'./stackix_work'

Source code in src/autoencodix/trainers/_stackix_orchestrator.py

def __init__(
    self,
    trainset: Optional[MultiModalDataset],
    validset: Optional[MultiModalDataset],
    config: DefaultConfig,
    model_type: Type[BaseAutoencoder],
    loss_type: Type[BaseLoss],
    testset: Optional[MultiModalDataset] = None,
    trainer_type: Type[GeneralTrainer] = GeneralTrainer,
    dataset_type: Type[BaseDataset] = NumericDataset,
    workdir: str = "./stackix_work",
):
    """
    Initialize the StackixOrchestrator with datasets and configuration.

    Args:
        trainset: Training dataset containing multiple modalities
        validset: Validation dataset containing multiple modalities
        config: Configuration parameters for training and model architecture
        model_type: Type of autoencoder model to use for each modality
        loss_type: Type of loss function to use for training
        testset: Dataset with test split
        trainer_type: Type to use for training (default is GeneralTrainer)
        dataset_type: Type to use for creating datasets (default is NumericDataset)
        workdir: Directory to save intermediate models and results (default is "./stackix_work")
    """
    self._trainset = trainset
    self._validset = validset
    self._testset = testset
    self._config = config
    self._model_type = model_type
    self._loss_type = loss_type
    self._workdir = workdir
    self._trainer_class = trainer_type
    self._dataset_type = dataset_type

    # Initialize fabric for distributed operations
    strategy = "auto"
    if hasattr(config, "gpu_strategy"):
        strategy = config.gpu_strategy

    self._fabric = lightning_fabric.Fabric(
        accelerator=(
            str(config.device) if hasattr(config, "device") else None
        ),  # ty: ignore[invalid-argument-type]
        devices=config.n_gpus if hasattr(config, "n_gpus") else 1,
        precision=(
            config.float_precision if hasattr(config, "float_precision") else "32"
        ),
        strategy=strategy,
    )

    # Initialize storage for models and results
    self.modality_trainers: Dict[str, BaseAutoencoder] = {}
    self.modality_results: Dict[str, Result] = {}
    self._modality_latent_dims: Dict[str, int] = {}
    self.concatenated_latent_spaces: Dict[str, torch.Tensor] = {}

    # Stacked model and trainer will be created during the process
    self.stacked_model: Optional[BaseAutoencoder] = None
    self.stacked_trainer: Optional[GeneralTrainer] = None
    self.concat_idx: Dict[str, Optional[Tuple[int, int]]] = {}
    self.dropped_indices_map: Dict[str, np.ndarray] = {}
    self.reconstruction_shapes: Dict[str, torch.Size] = {}
    self.common_sample_ids: Optional[pd.Index] = None

    # Create the directory if it doesn't exist
    os.makedirs(name=self._workdir, exist_ok=True)

predict_modalities(data)

Predicts using the trained models for each modality.

Parameters:

Name	Type	Description	Default
data	MultiModalDataset	Input data for prediction, uses test data if not provided	required

Returns:

Type	Description
Dict[str, Tensor]	Dictionary of reconstructed tensors by modality

Raises:

Type	Description
ValueError	If model has not been trained yet or no data is available

Source code in src/autoencodix/trainers/_stackix_orchestrator.py

def predict_modalities(self, data: MultiModalDataset) -> Dict[str, torch.Tensor]:
    """Predicts using the trained models for each modality.

    Args:
        data: Input data for prediction, uses test data if not provided

    Returns:
        Dictionary of reconstructed tensors by modality

    Raises:
        ValueError: If model has not been trained yet or no data is available
    """
    if not self.modality_trainers:
        raise ValueError(
            "No modality models have been trained. Call train_modalities() first."
        )

    predictions = {}
    for modality, trainer in self.modality_trainers.items():
        predictions[modality] = trainer.predict(
            data=data.datasets[modality],
            model=trainer._model,
        )
    return predictions

prepare_latent_datasets(split)

Prepares datasets with concatenated latent spaces for stacked model training.

This is the second phase of Stackix training where latent spaces from all modalities are extracted and concatenated.

Returns:

Type	Description
NumericDataset	Training and validation datasets with concatenated latent spaces

Raises:

Type	Description
ValueError	If no modality models have been trained or no latent spaces could be extracted

Source code in src/autoencodix/trainers/_stackix_orchestrator.py

def prepare_latent_datasets(self, split: str) -> NumericDataset:
    """Prepares datasets with concatenated latent spaces for stacked model training.

    This is the second phase of Stackix training where latent spaces from
    all modalities are extracted and concatenated.

    Returns:
        Training and validation datasets with concatenated latent spaces

    Raises:
        ValueError: If no modality models have been trained or no latent spaces could be extracted
    """
    if not self.modality_trainers:
        raise ValueError(
            "No modality models have been trained. Call train_modalities() first."
        )

    # Extract and concatenate latent spaces
    if split == "test":
        if self._testset is None:
            raise ValueError(
                "No test dataset available. Please provide a test dataset for prediction."
            )
        latent = self._extract_latent_spaces(
            result_dict=self.predict_modalities(data=self._testset), split=split
        )
        ds = NumericDataset(
            data=latent,
            config=self._config,
            sample_ids=self.common_sample_ids,
            feature_ids=[
                f"{k}_latent_{i}"
                for k, (start, end) in self.concat_idx.items()
                for i in range(start, end)
            ],
        )
        return ds

    latent = self._extract_latent_spaces(
        result_dict=self.modality_results, split=split
    )
    # Create datasets for the concatenated latent spaces
    feature_ids = [
        f"{k}_latent_{i}"
        for k, (start, end) in self.concat_idx.items()
        for i in range(start, end)
    ]
    ds = NumericDataset(
        data=latent,
        config=self._config,
        sample_ids=self.common_sample_ids,
        feature_ids=feature_ids,
    )

    return ds

reconstruct_from_stack(reconstructed_stack)

Reconstructs the full data for each modality from the stacked latent reconstruction.

Parameters:

Name	Type	Description	Default
reconstructed_stack	Tensor	Tensor with the reconstructed concatenated latent space	required

Returns: Dictionary of reconstructed tensors by modality

Source code in src/autoencodix/trainers/_stackix_orchestrator.py

def reconstruct_from_stack(
    self, reconstructed_stack: torch.Tensor
) -> Dict[str, torch.Tensor]:
    """Reconstructs the full data for each modality from the stacked latent reconstruction.

    Args:
        reconstructed_stack: Tensor with the reconstructed concatenated latent space
    Returns:
        Dictionary of reconstructed tensors by modality
    """
    modality_reconstructions: Dict[str, torch.Tensor] = {}

    for trainer in self.modality_trainers.values():
        trainer._model.eval()

    for name, (start_idx, end_idx) in self.concat_idx.items():
        # =================================================================
        # STEP 1: DE-CONCATENATE THE ALIGNED PART
        # This is the small, dense reconstruction matching the common samples.
        # =================================================================
        aligned_latent_recon = reconstructed_stack[:, start_idx:end_idx]

        # =================================================================
        # STEP 2: RE-ASSEMBLE THE FULL-SIZE LATENT SPACE
        # This is the logic from the old `reconstruct_full_latent` function,
        # now integrated directly here.
        # =================================================================
        original_shape = self.reconstruction_shapes[name]
        dropped_indices = self.dropped_indices_map[name]
        n_original_samples = original_shape[0]

        # Determine the original integer positions of the samples that were KEPT.
        kept_indices = np.delete(np.arange(n_original_samples), dropped_indices)

        # Create a full-size placeholder tensor matching the original latent space.
        # We use zeros, as they are a neutral input for most decoders.
        full_latent_recon = torch.zeros(
            size=original_shape,
            dtype=aligned_latent_recon.dtype,
            device=self._fabric.device,
        )

        # Place the aligned reconstruction data into the correct rows of the full tensor.
        full_latent_recon[kept_indices, :] = self._fabric.to_device(
            aligned_latent_recon
        )

        # =================================================================
        # STEP 3: DECODE THE FULL-SIZE LATENT TENSOR
        # The decoder now receives a tensor with the correct number of samples,
        # including the rows of zeros for the dropped samples.
        # =================================================================
        with torch.no_grad():
            model = self.modality_trainers[name].get_model()
            model = self._fabric.to_device(model)

            final_recon = model.decode(full_latent_recon)
            modality_reconstructions[name] = final_recon.cpu()

    return modality_reconstructions

set_testset(testset)

Set the test dataset for the orchestrator.

Parameters:

Name	Type	Description	Default
testset	MultiModalDataset	Test dataset containing multiple modalities	required

Source code in src/autoencodix/trainers/_stackix_orchestrator.py

def set_testset(self, testset: MultiModalDataset) -> None:
    """Set the test dataset for the orchestrator.

    Args:
        testset: Test dataset containing multiple modalities
    """
    self._testset = testset

train_modalities()

Trains all modality-specific models.

This is the first phase of Stackix training where each modality is trained independently before their latent spaces are combined.

Returns:

Type	Description
Dict[str, BaseAutoencoder]	Dictionary of trained models for each modality and Dictionary of training results
Dict[str, Result]	for each modality.

Raises:

Type	Description
ValueError	If trainset is not a MultiModalDataset or has no modalities

Source code in src/autoencodix/trainers/_stackix_orchestrator.py

def train_modalities(self) -> Tuple[Dict[str, BaseAutoencoder], Dict[str, Result]]:
    """Trains all modality-specific models.

    This is the first phase of Stackix training where each modality is
    trained independently before their latent spaces are combined.

    Returns:
        Dictionary of trained models for each modality and Dictionary of training results
        for each modality.

    Raises:
        ValueError: If trainset is not a MultiModalDataset or has no modalities
    """
    # if not isinstance(self._trainset, MulDataset):
    #     raise ValueError("Trainset must be a MultiModalDataset for Stackix training")

    keys = self._trainset.datasets.keys()
    if not keys:
        raise ValueError("No modalities found in trainset")

    n_modalities = len(keys)

    # Determine if we should use distributed training
    n_gpus = self._config.n_gpus if hasattr(self._config, "n_gpus") else 0
    use_distributed = bool(n_gpus and n_gpus > 1 and n_modalities > 1)

    if use_distributed:
        self._train_distributed(keys=keys)
    else:
        self._train_sequential(keys=keys)

    return self.modality_trainers, self.modality_results

StackixTrainer

Bases: GeneralTrainer

StackixTrainer is a wrapper for StackixOrchestrator that conforms to the BaseTrainer interface.

This trainer maintains compatibility with the BasePipeline interface while leveraging the more modular and well-designed StackixOrchestrator classes for the actual implementation.

Attributes: _workdir: Directory for saving intermediate models and results _result: Result object to store training outcomes _config: Configuration parameters for training and model architecture _model_type: Type of autoencoder model to use for each modality _loss_type: Type of loss function to use for training _trainset: Training dataset containing multiple modalities _validset: Validation dataset containing multiple modalities _orchestrator_type: Type to use for orchestrating modality training (default is StackixOrchestrator) _trainer_type: Type to use for training each modality model (default is GeneralTrainer) _modality_trainers: Dictionary of trained models for each modality _modality_results: Dictionary of training results for each modality _trainer: Trainer for the stacked model _fabric: Lightning Fabric wrapper for device and precision management _train_latent_ds: Training dataset with concatenated latent spaces _valid_latent_ds: Validation dataset with concatenated latent spaces concat_idx: Indices used for concatenating latent spaces _model: The instantiated stacked model architecture _optimizer: The optimizer used for training _orchestrator: The orchestrator that manages modality model training and latent space preparation _workdir: Directory for saving intermediate models and results

Source code in src/autoencodix/trainers/_stackix_trainer.py

class StackixTrainer(GeneralTrainer):
    """StackixTrainer is a wrapper for StackixOrchestrator that conforms to the BaseTrainer interface.

    This trainer maintains compatibility with the BasePipeline interface while
    leveraging the more modular and well-designed StackixOrchestrator
    classes for the actual implementation.

    Attributes:
    _workdir: Directory for saving intermediate models and results
    _result: Result object to store training outcomes
    _config: Configuration parameters for training and model architecture
    _model_type: Type of autoencoder model to use for each modality
    _loss_type: Type of loss function to use for training
    _trainset: Training dataset containing multiple modalities
    _validset: Validation dataset containing multiple modalities
    _orchestrator_type: Type to use for orchestrating modality training (default is StackixOrchestrator)
    _trainer_type: Type to use for training each modality model (default is GeneralTrainer)
    _modality_trainers: Dictionary of trained models for each modality
    _modality_results: Dictionary of training results for each modality
    _trainer: Trainer for the stacked model
    _fabric: Lightning Fabric wrapper for device and precision management
    _train_latent_ds: Training dataset with concatenated latent spaces
    _valid_latent_ds: Validation dataset with concatenated latent spaces
    concat_idx: Indices used for concatenating latent spaces
    _model: The instantiated stacked model architecture
    _optimizer: The optimizer used for training
    _orchestrator: The orchestrator that manages modality model training and latent space preparation
    _workdir: Directory for saving intermediate models and results
    """

    def __init__(
        self,
        trainset: Optional[StackixDataset],
        validset: Optional[StackixDataset],
        result: Result,
        config: DefaultConfig,
        model_type: Type[BaseAutoencoder],
        loss_type: Type[BaseLoss],
        orchestrator_type: Type[StackixOrchestrator] = StackixOrchestrator,
        trainer_type: Type[BaseTrainer] = GeneralTrainer,
        workdir: str = "./stackix_work",
        ontologies: Optional[Tuple] = None,
        **kwargs,
    ) -> None:
        """Initialize the StackixTrainer with datasets and configuration.

        Args:
            trainset: Training dataset containing multiple modalities
            validset: Validation dataset containing multiple modalities
            result: Result object to store training outcomes
            config: Configuration parameters for training and model architecture
            model_type: Type of autoencoder model to use for each modality
            loss_type: Type of loss function to use for training
            orchestrator_type: Type to use for orchestrating modality training (default is StackixOrchestrator)
            trainer_type: Type to use for training each modality model (default is GeneralTrainer)
            workdir: Directory to save intermediate models and results (default is "./stackix_work")
            onotologies: Ontology information, if provided for Ontix compatibility
        """
        self._workdir = workdir
        self._result = result
        self._config = config
        self._model_type = model_type
        self._loss_type = loss_type
        self._trainset = trainset
        self._validset = validset
        self._orchestrator_type = orchestrator_type
        self._trainer_type = trainer_type

        self._orchestrator = self._orchestrator_type(
            trainset=trainset,
            validset=validset,
            config=config,
            model_type=model_type,
            loss_type=loss_type,
            workdir=workdir,
        )
        self._modality_trainers: Optional[Dict[str, BaseAutoencoder]] = None
        self._modality_results: Optional[Dict[str, Result]] = None

        self._fabric = Fabric(
            accelerator=self._config.device,
            devices=self._config.n_gpus,
            precision=self._config.float_precision,
            strategy=self._config.gpu_strategy,
        )

        super().__init__(
            trainset=trainset,
            validset=validset,
            result=result,
            config=config,
            model_type=model_type,
            loss_type=loss_type,
        )

    def get_model(self) -> torch.nn.Module:
        """Getter for the the trained model.

        Returns:
            The trained model
        """
        return self._model

    def train(self) -> Result:
        """Train the stacked model on the concatenated latent space.

        Uses the standard BaseTrainer training process but with the stacked model.

        Returns:
            Training results including losses, latent spaces, and other metrics
        """
        print("Training each modality model...")
        self._train_modalities()
        print("finished training each modality model")
        self._trainer = self._trainer_type(
            trainset=self._train_latent_ds,
            validset=self._valid_latent_ds,
            result=self._result,
            config=self._config,
            model_type=self._model_type,
            loss_type=self._loss_type,
        )

        self._model = self._trainer._model
        self._result = self._trainer.train()

        return self._result

    def _train_modalities(self) -> None:
        """Trains a Autoencoder for each modality in the dataset.
        This method orchestrates the training of individual modality models

        This method orchestrates the complete training process:
        1. Train individual modality models
        2. Extract and concatenate latent spaces
        3. Populates the self._modality_models and self._modality_results attributes

        """
        # Step 1: Train individual modality models
        self._modality_trainers, self._modality_results = (
            self._orchestrator.train_modalities()
        )

        self._result.sub_results = self._modality_results
        # Step 2: Prepare concatenated latent space datasets
        self._train_latent_ds = self._orchestrator.prepare_latent_datasets(
            split="train"
        )
        self._valid_latent_ds = self._orchestrator.prepare_latent_datasets(
            split="valid"
        )
        self.concat_idx = self._orchestrator.concat_idx

    def _reconstruct(self, split: str) -> None:
        """Orchestrates the reconstruction by delegating the task to the StackixOrchestrator.

        Args:
            split: The data split to reconstruct ('train', 'valid', 'test').
        """
        stacked_recon = self._result.reconstructions.get(epoch=-1, split=split)
        if stacked_recon is None:
            print(f"Warning: No reconstruction found for split '{split}'. Skipping.")
            return

        # The orchestrator handles the de-concatenation, re-assembly, and decoding.
        modality_reconstructions = self._orchestrator.reconstruct_from_stack(
            reconstructed_stack=torch.from_numpy(stacked_recon).to(self._fabric.device)
        )

        # The result is the final, full-sized data reconstructions.
        self._result.sub_reconstructions = modality_reconstructions

    def predict(
        self, data: BaseDataset, model: Optional[torch.nn.Module] = None, **kwargs
    ) -> Result:
        """Make predictions on the given dataset.

        Args:
            data: The dataset to make predictions on.
            model: The model to use for predictions. If None, uses the trained model.
            **kwargs: Additional keyword arguments.
        Returns:
            Result: The prediction results including reconstructions and latent spaces.
        """
        self.n_test = len(data) if data is not None else 0
        self._orchestrator.set_testset(testset=data)
        test_ds = self._orchestrator.prepare_latent_datasets(split="test")
        self.testset = test_ds
        print(test_ds)
        pred_result = super().predict(data=test_ds, model=model)
        self._result.update(other=pred_result)
        self._reconstruct(split="test")
        return self._result

__init__(trainset, validset, result, config, model_type, loss_type, orchestrator_type=StackixOrchestrator, trainer_type=GeneralTrainer, workdir='./stackix_work', ontologies=None, **kwargs)

Initialize the StackixTrainer with datasets and configuration.

Parameters:

Name	Type	Description	Default
trainset	Optional[StackixDataset]	Training dataset containing multiple modalities	required
validset	Optional[StackixDataset]	Validation dataset containing multiple modalities	required
result	Result	Result object to store training outcomes	required
config	DefaultConfig	Configuration parameters for training and model architecture	required
model_type	Type[BaseAutoencoder]	Type of autoencoder model to use for each modality	required
loss_type	Type[BaseLoss]	Type of loss function to use for training	required
orchestrator_type	Type[StackixOrchestrator]	Type to use for orchestrating modality training (default is StackixOrchestrator)	StackixOrchestrator
trainer_type	Type[BaseTrainer]	Type to use for training each modality model (default is GeneralTrainer)	GeneralTrainer
workdir	str	Directory to save intermediate models and results (default is "./stackix_work")	'./stackix_work'
onotologies		Ontology information, if provided for Ontix compatibility	required

Source code in src/autoencodix/trainers/_stackix_trainer.py

def __init__(
    self,
    trainset: Optional[StackixDataset],
    validset: Optional[StackixDataset],
    result: Result,
    config: DefaultConfig,
    model_type: Type[BaseAutoencoder],
    loss_type: Type[BaseLoss],
    orchestrator_type: Type[StackixOrchestrator] = StackixOrchestrator,
    trainer_type: Type[BaseTrainer] = GeneralTrainer,
    workdir: str = "./stackix_work",
    ontologies: Optional[Tuple] = None,
    **kwargs,
) -> None:
    """Initialize the StackixTrainer with datasets and configuration.

    Args:
        trainset: Training dataset containing multiple modalities
        validset: Validation dataset containing multiple modalities
        result: Result object to store training outcomes
        config: Configuration parameters for training and model architecture
        model_type: Type of autoencoder model to use for each modality
        loss_type: Type of loss function to use for training
        orchestrator_type: Type to use for orchestrating modality training (default is StackixOrchestrator)
        trainer_type: Type to use for training each modality model (default is GeneralTrainer)
        workdir: Directory to save intermediate models and results (default is "./stackix_work")
        onotologies: Ontology information, if provided for Ontix compatibility
    """
    self._workdir = workdir
    self._result = result
    self._config = config
    self._model_type = model_type
    self._loss_type = loss_type
    self._trainset = trainset
    self._validset = validset
    self._orchestrator_type = orchestrator_type
    self._trainer_type = trainer_type

    self._orchestrator = self._orchestrator_type(
        trainset=trainset,
        validset=validset,
        config=config,
        model_type=model_type,
        loss_type=loss_type,
        workdir=workdir,
    )
    self._modality_trainers: Optional[Dict[str, BaseAutoencoder]] = None
    self._modality_results: Optional[Dict[str, Result]] = None

    self._fabric = Fabric(
        accelerator=self._config.device,
        devices=self._config.n_gpus,
        precision=self._config.float_precision,
        strategy=self._config.gpu_strategy,
    )

    super().__init__(
        trainset=trainset,
        validset=validset,
        result=result,
        config=config,
        model_type=model_type,
        loss_type=loss_type,
    )

get_model()

Getter for the the trained model.

Returns:

Type	Description
Module	The trained model

Source code in src/autoencodix/trainers/_stackix_trainer.py

def get_model(self) -> torch.nn.Module:
    """Getter for the the trained model.

    Returns:
        The trained model
    """
    return self._model

predict(data, model=None, **kwargs)

Make predictions on the given dataset.

Parameters:

Name	Type	Description	Default
data	BaseDataset	The dataset to make predictions on.	required
model	Optional[Module]	The model to use for predictions. If None, uses the trained model.	None
**kwargs		Additional keyword arguments.	{}

Returns: Result: The prediction results including reconstructions and latent spaces.

Source code in src/autoencodix/trainers/_stackix_trainer.py

def predict(
    self, data: BaseDataset, model: Optional[torch.nn.Module] = None, **kwargs
) -> Result:
    """Make predictions on the given dataset.

    Args:
        data: The dataset to make predictions on.
        model: The model to use for predictions. If None, uses the trained model.
        **kwargs: Additional keyword arguments.
    Returns:
        Result: The prediction results including reconstructions and latent spaces.
    """
    self.n_test = len(data) if data is not None else 0
    self._orchestrator.set_testset(testset=data)
    test_ds = self._orchestrator.prepare_latent_datasets(split="test")
    self.testset = test_ds
    print(test_ds)
    pred_result = super().predict(data=test_ds, model=model)
    self._result.update(other=pred_result)
    self._reconstruct(split="test")
    return self._result

train()

Train the stacked model on the concatenated latent space.

Uses the standard BaseTrainer training process but with the stacked model.

Returns:

Type	Description
Result	Training results including losses, latent spaces, and other metrics

Source code in src/autoencodix/trainers/_stackix_trainer.py

def train(self) -> Result:
    """Train the stacked model on the concatenated latent space.

    Uses the standard BaseTrainer training process but with the stacked model.

    Returns:
        Training results including losses, latent spaces, and other metrics
    """
    print("Training each modality model...")
    self._train_modalities()
    print("finished training each modality model")
    self._trainer = self._trainer_type(
        trainset=self._train_latent_ds,
        validset=self._valid_latent_ds,
        result=self._result,
        config=self._config,
        model_type=self._model_type,
        loss_type=self._loss_type,
    )

    self._model = self._trainer._model
    self._result = self._trainer.train()

    return self._result

XModalTrainer

Bases: BaseTrainer

Trainer for cross-modal autoencoders, implements multimodal training with adversarial component.

Attributes:

Name	Type	Description
_trainset		The dataset used for training, must be a MultiModalDataset.
_n_train_modalities		Number of modalities in the training dataset.
model_map		Mapping from DataSetTypes to specific autoencoder architectures.
model_trainer_map		Mapping from autoencoder architectures to their corresponding trainer classes.
_n_cpus		Number of CPU cores available for data loading.
latent_dim		Dimensionality of the shared latent space.
n_test	Optional[int]	Number of samples in the test set, if provided.
n_train		Number of samples in the training set.
n_valid		Number of samples in the validation set, if provided.
n_features		Total number of features across all modalities.
_cur_epoch	int	Current epoch number during training.
_is_checkpoint_epoch	Optional[bool]	Flag indicating if the current epoch is a checkpoint epoch.
_sub_loss_type		Loss function type for individual modality autoencoders.
sub_loss		Instantiated loss function for individual modality autoencoders.
_clf_epoch_loss	float	Cumulative classifier loss for the current epoch.
_epoch_loss	float	Cumulative total loss for the current epoch.
_epoch_loss_valid	float	Cumulative total validation loss for the current epoch.
_modality_dynamics	float	Dictionary holding model, optimizer, and training state for each modality.
_latent_clf	float	Classifier model for adversarial training on latent spaces.
_clf_optim	float	Optimizer for the classifier model.
_clf_loss_fn	float	Loss function for the classifier.
_trainloader	float	DataLoader for the training dataset.
_validloader	float	DataLoader for the validation dataset, if provided.
_model	float	The instantiated stacked model architecture.
_validset		The dataset used for validation, if provided, must be a MultiModalDataset.
_fabric		Lightning Fabric wrapper for device and precision management.

Source code in src/autoencodix/trainers/_xmodal_trainer.py

class XModalTrainer(BaseTrainer):
    """Trainer for cross-modal autoencoders, implements multimodal training with adversarial component.

    Attributes:
        _trainset: The dataset used for training, must be a MultiModalDataset.
        _n_train_modalities: Number of modalities in the training dataset.
        model_map: Mapping from DataSetTypes to specific autoencoder architectures.
        model_trainer_map: Mapping from autoencoder architectures to their corresponding trainer classes.
        _n_cpus: Number of CPU cores available for data loading.
        latent_dim: Dimensionality of the shared latent space.
        n_test: Number of samples in the test set, if provided.
        n_train: Number of samples in the training set.
        n_valid: Number of samples in the validation set, if provided.
        n_features: Total number of features across all modalities.
        _cur_epoch: Current epoch number during training.
        _is_checkpoint_epoch: Flag indicating if the current epoch is a checkpoint epoch.
        _sub_loss_type: Loss function type for individual modality autoencoders.
        sub_loss: Instantiated loss function for individual modality autoencoders.
        _clf_epoch_loss: Cumulative classifier loss for the current epoch.
        _epoch_loss: Cumulative total loss for the current epoch.
        _epoch_loss_valid: Cumulative total validation loss for the current epoch.
        _modality_dynamics: Dictionary holding model, optimizer, and training state for each modality.
        _latent_clf: Classifier model for adversarial training on latent spaces.
        _clf_optim: Optimizer for the classifier model.
        _clf_loss_fn: Loss function for the classifier.
        _trainloader: DataLoader for the training dataset.
        _validloader: DataLoader for the validation dataset, if provided.
        _model: The instantiated stacked model architecture.
        _validset: The dataset used for validation, if provided, must be a MultiModalDataset.
        _fabric: Lightning Fabric wrapper for device and precision management.

    """

    def __init__(
        self,
        trainset: MultiModalDataset,
        validset: MultiModalDataset,
        result: Result,
        config: DefaultConfig,
        model_type: Type[BaseAutoencoder],
        loss_type: Type[BaseLoss],
        model_map: Dict[DataSetTypes, Type[BaseAutoencoder]],
        sub_loss_type: Type[BaseLoss] = VarixLoss,
        ontologies: Optional[Union[Tuple, List]] = None,
        **kwargs,
    ):
        """Initializes the XModalTrainer with datasets and configuration.

        Args:
            trainset: Training dataset containing multiple modalities
            validset: Validation dataset containing multiple modalities
            result: Result object to store training outcomes
            config: Configuration parameters for training and model architecture
            model_type: Type of autoencoder model to use for each modality (not used directly)
            loss_type: Type of loss function to use for training the stacked model
            sub_loss_type: Type of loss function to use for individual modality autoencoders
            model_map: Mapping from DataSetTypes to specific autoencoder architectures
            ontologies: Ontology information, for compatibility with Ontix
        """
        self._trainset = trainset
        self._n_train_modalities = self._trainset.n_modalities  # type: ignore
        self.model_map = model_map
        self.model_trainer_map = model_trainer_map
        self._n_cpus = os.cpu_count()
        if self._n_cpus is None:
            self._n_cpus = 0

        super().__init__(
            trainset, validset, result, config, model_type, loss_type, ontologies
        )
        # init attributes ------------------------------------------------
        self.latent_dim = config.latent_dim
        self.n_test: Optional[int] = None
        self.n_train = len(trainset.data) if trainset else 0
        self.n_valid = len(validset.data) if validset else 0
        self._cur_epoch: int = 0
        self._is_checkpoint_epoch: Optional[bool] = None
        self._sub_loss_type = sub_loss_type
        self.sub_loss = sub_loss_type(config=self._config)
        self._clf_epoch_loss: float = 0.0
        self._epoch_loss: float = 0.0
        self._epoch_loss_valid: float = 0.0
        self._all_dynamics: Dict[str, Dict[int, Dict]] = {
            "test": defaultdict(dict),
            "train": defaultdict(dict),
            "valid": defaultdict(dict),
        }

    def _init_model_architecture(self, ontologies: Optional[Tuple] = None):
        """Override parent's model init - we don't use a single model, needs to be there because is abstract in parent."""
        pass

    def _setup_fabric(self, old_model=None):
        """Sets up the models, optimizers, and data loaders with Lightning Fabric."""

        self._fabric.launch()
        self._init_adversarial_training()
        self._init_modality_training()
        self._latent_clf, self._clf_optim = self._fabric.setup(
            self._latent_clf, self._clf_optim
        )

        self._trainloader = self._fabric.setup_dataloaders(self._trainloader)
        if self._validloader is not None:
            self._validloader = self._fabric.setup_dataloaders(self._validloader)

        for dynamics in self._modality_dynamics.values():
            dynamics["model"], dynamics["optim"] = self._fabric.setup(
                dynamics["model"], dynamics["optim"]
            )
            dynamics["model"].mark_forward_method("decode")

    def _init_loaders(self):
        """Initializes DataLoaders with smart sampler selection based on pairing."""

        def _build_loader(dataset: MultiModalDataset, is_train: bool):
            batch_size = self._config.batch_size
            collate_fn = create_multimodal_collate_fn(multimodal_dataset=dataset)

            if dataset.is_fully_paired:
                return DataLoader(
                    dataset,
                    batch_size=batch_size,
                    shuffle=is_train,
                    drop_last=is_train,
                    collate_fn=collate_fn,
                    pin_memory=self._config.pin_memory,
                )
            else:
                print(
                    f"Dataset has UNPAIRED samples → using CoverageEnsuringSampler "
                    f"({len(dataset.paired_sample_ids)} paired + {len(dataset.unpaired_sample_ids)} unpaired)"
                )
                sampler = CoverageEnsuringSampler(
                    multimodal_dataset=dataset,
                    batch_size=batch_size,
                )
                return DataLoader(
                    dataset,
                    batch_sampler=sampler,  # note: batch_sampler, not sampler
                    collate_fn=collate_fn,
                    pin_memory=self._config.pin_memory,
                )

        # Build train and validation loaders
        self._trainloader = _build_loader(self._trainset, is_train=True)
        self._validloader = _build_loader(self._validset, is_train=False)

        # Optional: expose for logging
        self._train_is_fully_paired = len(self._trainset.unpaired_sample_ids) == 0

    def _init_modality_training(self):
        """Initializes models, optimizers, and training state for each modality."""
        self._modality_dynamics = {
            mod_name: None for mod_name in self._trainset.datasets.keys()
        }

        self._result.sub_results = {
            mod_name: None for mod_name in self._trainset.datasets.keys()
        }

        data_info = self._config.data_config.data_info
        for mod_name, ds in self._trainset.datasets.items():
            simple_name = mod_name.split(".")[1]
            local_epochs = data_info[simple_name].pretrain_epochs

            pretrain_epochs = (
                local_epochs
                if local_epochs is not None
                else self._config.pretrain_epochs
            )

            model_type = self.model_map.get(ds.mytype)
            if model_type is None:
                raise ValueError(
                    f"No Mapping exists for {ds.mytype}, you passed this mapping: {self.model_map}"
                )
            model = model_type(config=self._config, input_dim=ds.get_input_dim())
            if (
                Version(torch.__version__) >= Version("2.0")
                and torch.cuda.is_available()
                and self._config.compile_model
            ):
                model = torch.compile(model)
            optimizer = torch.optim.AdamW(
                params=model.parameters(),
                lr=self._config.learning_rate,
                weight_decay=self._config.weight_decay,
            )
            self._modality_dynamics[mod_name] = {
                "model": model,
                "optim": optimizer,
                "mp": [],
                "losses": [],
                "config_name": simple_name,
                "pretrain_epochs": pretrain_epochs,
                "mytype": ds.mytype,
                "pretrain_result": None,
            }

    def _init_adversarial_training(self):
        """Initializes the classifier and its optimizer for adversarial training."""
        self._latent_clf = Classifier(
            input_dim=self._config.latent_dim, n_modalities=self._n_train_modalities
        )
        self._clf_optim = torch.optim.AdamW(
            params=self._latent_clf.parameters(),
            lr=self._config.learning_rate,
            weight_decay=self._config.weight_decay,
        )
        self._clf_loss_fn = torch.nn.CrossEntropyLoss(
            reduction=self._config.loss_reduction
        )

    def _modalities_forward(self, batch: Dict[str, Dict[str, Any]]):
        """Performs forward pass for each modality in the batch and computes losses."""
        for k, v in batch.items():
            model = self._modality_dynamics[k]["model"]
            data = v["data"]

            mp = model(data)

            loss, loss_stats = self.sub_loss(
                model_output=mp, targets=data, epoch=self._cur_epoch
            )

            self._modality_dynamics[k]["loss_stats"] = loss_stats
            self._modality_dynamics[k]["loss"] = loss
            self._modality_dynamics[k]["mp"] = mp

    def _prep_adver_training(self) -> Tuple[torch.Tensor, torch.Tensor]:
        """Prepares concatenated latent spaces and corresponding labels for adversarial training.

        Returns:
            A tuple containing:
            - latents: Concatenated latent space tensor of shape (total_samples, latent_dim)
            - labels: Tensor of modality labels of shape (total_samples,)
        """
        all_latents: List[torch.Tensor] = []
        all_labels: List[torch.Tensor] = []
        mod2idx = {
            mod_name: idx for idx, mod_name in enumerate(self._modality_dynamics.keys())
        }

        for mod_name, helper in self._modality_dynamics.items():
            output: ModelOutput = helper["mp"]
            latents = output.latentspace
            label_id = mod2idx[mod_name]
            all_latents.append(latents)
            all_labels.append(
                torch.full(
                    (latents.size(0),),
                    fill_value=label_id,
                    dtype=torch.long,
                    device=self._fabric.device,
                )
            )

        return torch.cat(all_latents, dim=0), torch.cat(all_labels, dim=0)

    def _train_clf(self, latents: torch.Tensor, labels: torch.Tensor) -> None:
        """Performs a single optimization step for the classifier.
        Args:
            latents: Concatenated latent space tensor of shape (total_samples, latent_dim)
            labels: Tensor of modality labels of shape (total_samples,)

        """
        self._clf_optim.zero_grad()
        # We must detach the latents from the computation graph. This serves two purposes:
        # 1. It isolates the classifier, ensuring that gradients only update its weights.
        # 2. It prevents a "two backward passes" error by leaving the graph to the
        #    autoencoders untouched, ready for use in the next stage.
        detached_latents = latents.detach()
        clf_scores = self._latent_clf(detached_latents)
        clf_loss = self._clf_loss_fn(clf_scores, labels)
        self._fabric.backward(clf_loss)
        self._clf_optim.step()
        self._clf_epoch_loss += clf_loss.item()

    def _validate_one_epoch(self) -> Tuple[List[Dict], Dict[str, float], int]:
        """Runs a single, combined validation pass for the entire model, including dedicated metrics for the classifier.

        Returns:
            A tuple containing:
            - epoch_dynamics: List of dictionaries capturing dynamics for each batch
            - sub_losses: Dictionary of accumulated sub-losses for the epoch
            - n_samples_total: Total number of samples processed in validation
        """
        for dynamics in self._modality_dynamics.values():
            dynamics["model"].eval()
        self._latent_clf.eval()

        self._epoch_loss_valid = 0.0
        total_clf_loss = 0.0
        n_samples_total: int = 0
        epoch_dynamics: List[Dict] = []
        sub_losses: Dict[str, float] = defaultdict(float)

        with torch.no_grad():
            for batch in self._validloader:
                with self._fabric.autocast():
                    self._modalities_forward(batch=batch)
                    latents, labels = self._prep_adver_training()
                    n_samples_total += len(labels)
                    clf_scores = self._latent_clf(latents)

                    batch_loss, loss_dict = self._loss_fn(
                        batch=batch,
                        modality_dynamics=self._modality_dynamics,
                        clf_scores=clf_scores,
                        labels=labels,
                        clf_loss_fn=self._clf_loss_fn,
                        is_training=False,
                    )
                    self._epoch_loss_valid += batch_loss.item()
                    for k, v in loss_dict.items():
                        value = v.item() if hasattr(v, "item") else v
                        if "_factor" not in k:
                            sub_losses[k] += value
                        else:
                            sub_losses[k] = value

                    clf_loss = self._clf_loss_fn(clf_scores, labels)
                    total_clf_loss += clf_loss.item()
                    if self._is_checkpoint_epoch:
                        batch_capture = self._capture_dynamics(batch)
                        epoch_dynamics.append(batch_capture)

        sub_losses["clf_loss"] = total_clf_loss
        n_samples_total /= (
            self._n_train_modalities
        )  # n_modalities because each sample is counted once per modality and n_modalites is the same for train and valid
        for k, v in sub_losses.items():
            if "_factor" not in k:
                sub_losses[k] = v / n_samples_total  # Average over all samples
        self._epoch_loss_valid = (
            self._epoch_loss_valid / n_samples_total
        )  # Average over all samples
        return epoch_dynamics, sub_losses, len(self._validset)

    def _train_one_epoch(self) -> Tuple[List[Dict], Dict[str, float], int]:
        """Runs the training loop with corrected adversarial logic.
        Returns:
            A tuple containing:
            - epoch_dynamics: List of dictionaries capturing dynamics for each batch
            - sub_losses: Dictionary of accumulated sub-losses for the epoch
            - n_samples_total: Total number of samples processed in training

        """
        for dynamics in self._modality_dynamics.values():
            dynamics["model"].train()
        self._latent_clf.train()

        self._clf_epoch_loss = 0
        self._epoch_loss = 0
        epoch_dynamics: List[Dict] = []
        sub_losses: Dict[str, float] = defaultdict(float)
        n_samples_total: int = (
            0  # because of unpaired training we need to sum the samples instead of using len(dataset)
        )

        for batch in self._trainloader:
            with self._fabric.autocast():
                # --- Stage 1: forward for each data modality ---
                self._modalities_forward(batch=batch)

                # --- Stage 2: Train the Classifier ---
                latents, labels = self._prep_adver_training()
                n_samples_total += latents.size(0)
                self._train_clf(latents=latents, labels=labels)

                # --- Stage 3: Train the Autoencoders ---
                for _, dynamics in self._modality_dynamics.items():
                    dynamics["optim"].zero_grad()

                # We re-calculate scores here and cannot reuse the `clf_scores` from Stage 2.
                # The previous scores were from detached latents and would block the gradients
                # needed to train the autoencoders adversarially.
                clf_scores_for_adv = self._latent_clf(latents)

                batch_loss, loss_dict = self._loss_fn(
                    batch=batch,
                    modality_dynamics=self._modality_dynamics,
                    clf_scores=clf_scores_for_adv,
                    labels=labels,
                    clf_loss_fn=self._clf_loss_fn,
                    is_training=True,
                )
            self._fabric.backward(batch_loss)
            for _, dynamics in self._modality_dynamics.items():
                dynamics["optim"].step()

            # --- Logging and Capturing ---
            self._epoch_loss += batch_loss.item()
            for k, v in loss_dict.items():
                value_to_add = v.item() if hasattr(v, "item") else v
                if "_factor" not in k:
                    sub_losses[k] += value_to_add
                else:
                    sub_losses[k] = value_to_add

            if self._is_checkpoint_epoch:
                batch_capture = self._capture_dynamics(batch)
                epoch_dynamics.append(batch_capture)
        n_samples_total /= self._n_train_modalities
        sub_losses["clf_loss"] = self._clf_epoch_loss
        for k, v in sub_losses.items():
            if "_factor" not in k:
                sub_losses[k] = v / n_samples_total  # Average over all samples
        self._epoch_loss = (
            self._epoch_loss / n_samples_total
        )  # Average over all samples

        return epoch_dynamics, sub_losses, n_samples_total

    def _pretraining(self):
        """Pretrain each modality's model if needed."""
        for mod_name, dynamic in self._modality_dynamics.items():
            mytype = dynamic.get("mytype")
            pretrain_epochs = dynamic.get("pretrain_epochs")
            print(f"Check if we need to pretrain: {mod_name}")
            print(f"pretrain epochs : {pretrain_epochs}")
            if not pretrain_epochs:
                print(f"No pretraining for {mod_name}")
                continue
            model_type = self.model_map.get(mytype)
            pretrainer_type = self.model_trainer_map.get(model_type)
            print(f"starting pretraining for: {mod_name} with {pretrainer_type}")
            trainset = self._trainset.datasets.get(mod_name)
            validset = self._validset.datasets.get(mod_name)
            pretrainer = pretrainer_type(
                trainset=trainset,
                validset=validset,
                result=Result(),
                config=self._config,
                model_type=model_type,
                loss_type=self._sub_loss_type,
                ontologies=self.ontologies,
            )
            pretrain_result = pretrainer.train(epochs_overwrite=pretrain_epochs)
            self._result.sub_results[f"pretrain.{mod_name}"] = pretrain_result
            self._modality_dynamics[mod_name]["pretrain_result"] = pretrain_result
            self._modality_dynamics[mod_name]["model"] = pretrain_result.model

    def train(self):
        """Orchestrates the full training process for the cross-modal autoencoder."""
        self._pretraining()
        for epoch in range(self._config.epochs):
            self._cur_epoch = epoch
            self._is_checkpoint_epoch = self._should_checkpoint(epoch=epoch)
            self._fabric.print(f"--- Epoch {epoch + 1}/{self._config.epochs} ---")
            if epoch == 0 and self._config.profiling:
                print("Profiling enabled for epoch 0")
                self._train_one_epoch_with_profiling()
                continue
            train_epoch_dynamics, train_sub_losses, n_samples_train = (
                self._train_one_epoch()
            )

            self._log_losses(
                split="train",
                total_loss=self._epoch_loss,
                sub_losses=train_sub_losses,
                n_samples=n_samples_train,
            )

            if self._validset:
                valid_epoch_dynamics, valid_sub_losses, n_samples_valid = (
                    self._validate_one_epoch()
                )
                self._log_losses(
                    split="valid",
                    total_loss=self._epoch_loss_valid,
                    sub_losses=valid_sub_losses,
                    n_samples=n_samples_valid,
                )
            if self._is_checkpoint_epoch:
                self._fabric.print(f"Storing checkpoint for epoch {epoch}...")
                self._store_checkpoint(
                    split="train", epoch_dynamics=train_epoch_dynamics
                )
                if self._validset:
                    self._store_checkpoint(
                        split="valid",
                        epoch_dynamics=valid_epoch_dynamics,
                    )
        # save final model
        self._store_checkpoint(split="train", epoch_dynamics=train_epoch_dynamics)
        self._store_final_models()
        return self._result

    def decode(self, x: torch.Tensor):
        """Decodes input latent representations
        Args:
            x: Latent representations to decode, shape (n_samples, latent_dim)
        """
        raise NotImplementedError("General decode step is not implemented for XModalix")

    def predict(
        self,
        data: BaseDataset,
        model: Optional[Dict[str, torch.nn.Module]] = None,
        **kwargs,
    ) -> Result:
        """Performs cross-modal prediction from a specified 'from' modality to a 'to' modality.

        The direction is determined by the 'translate_direction' attribute in the config.
        Results are stored in the Result object under split='test' and epoch=-1.

        Args:
            data: A MultiModalDataset containing the input data for the 'from' modality.

        Returns:
            The Result object populated with prediction results.
        """
        split: str = kwargs.pop("split", "test")
        if split not in ["train", "valid", "test"]:
            raise ValueError(
                f"split must be one of 'train', 'valid', or 'test', got {split}"
            )
        if not isinstance(data, MultiModalDataset):
            raise TypeError(
                f"type of data has to be MultiModalDataset, got: {type(data)}"
            )
        if model is not None and not hasattr(self, "_modality_dynamics"):
            self._modality_dynamics = {mod_name: {} for mod_name in model.keys()}

            for mod_name, mod_model in model.items():
                mod_model.eval()
                if not isinstance(mod_model, _FabricModule):
                    mod_model = self._fabric.setup_module(mod_model)
                mod_model.to(self._fabric.device)
                self._modality_dynamics[mod_name]["model"] = mod_model
            # self._setup_fabric(old_model=model)

        from_key = kwargs.pop("from_key", None)
        to_key = kwargs.pop("to_key", None)
        predict_keys = find_translation_keys(
            config=self._config,
            trained_modalities=list(self._modality_dynamics.keys()),
            from_key=from_key,
            to_key=to_key,
        )
        from_key, to_key = predict_keys["from"], predict_keys["to"]
        from_modality = self._modality_dynamics[from_key]
        to_modality = self._modality_dynamics[to_key]
        from_model = from_modality["model"]
        to_model = to_modality["model"]
        from_model.eval(), to_model.eval()
        # drop_last handled in custom sampler
        inference_loader = DataLoader(
            data,
            batch_sampler=CoverageEnsuringSampler(
                multimodal_dataset=data, batch_size=self._config.batch_size
            ),
            collate_fn=create_multimodal_collate_fn(multimodal_dataset=data),
        )
        inference_loader = self._fabric.setup_dataloaders(inference_loader)  # type: ignore
        epoch_dynamics: List[Dict] = []
        with (
            self._fabric.autocast(),
            torch.inference_mode(),
        ):
            for batch in inference_loader:
                # needed for visualize later
                self._get_vis_dynamics(batch=batch)
                from_z = self._modality_dynamics[from_key]["mp"].latentspace
                translated = to_model.decode(x=from_z)

                # for reference
                to_z = self._modality_dynamics[to_key]["mp"].latentspace
                to_to_reference = to_model.decode(x=to_z)

                batch_capture = self._capture_dynamics(batch)
                translation_key = "translation"

                reference_key = f"reference_{to_key}_to_{to_key}"
                batch_capture["reconstructions"][
                    translation_key
                ] = translated.cpu().numpy()

                batch_capture["reconstructions"][
                    reference_key
                ] = to_to_reference.cpu().numpy()

                if "sample_ids" in batch[from_key]:
                    batch_capture["sample_ids"][translation_key] = np.array(
                        batch[from_key]["sample_ids"]
                    )
                if "sample_ids" in batch[to_key]:
                    batch_capture["sample_ids"][reference_key] = np.array(
                        batch[to_key]["sample_ids"]
                    )

                epoch_dynamics.append(batch_capture)

        self._dynamics_to_result(split=split, epoch_dynamics=epoch_dynamics)

        print("Prediction complete.")
        return self._result

    def _get_vis_dynamics(self, batch: Dict[str, Dict[str, Any]]):
        """Runs a forward pass for each modality in the batch and stores the model outputs.

        This is used to capture dynamics for visualization or analysis not for actual translation as specified in predict.
        Args:
            batch: A dictionary where keys are modality names and values are dictionaries containing 'data' tensors.
        """
        for mod_name, mod_data in batch.items():
            model = self._modality_dynamics[mod_name]["model"]
            # Store model output (containing latents, recons, etc.)
            self._modality_dynamics[mod_name]["mp"] = model(mod_data["data"])

    def _capture_dynamics(
        self,
        batch_data: Union[Dict[str, Dict[str, Any]], Any],
    ) -> Union[Dict[str, Dict[str, np.ndarray]], Any]:
        """Captures and returns the dynamics (latents, reconstructions, etc.) for each modality in the batch.

        Args:
            batch_data: A dictionary where keys are modality names and values are dictionaries containing 'data' tensors
                         and optionally 'sample_ids'.
        Returns:
            A dictionary with keys 'latentspaces', 'reconstructions', 'mus', 'sigmas', and 'sample_ids',
            each mapping to another dictionary where keys are modality names and values are numpy arrays.
        """
        captured_data: Dict[str, Dict[str, Any]] = {
            "latentspaces": {},
            "reconstructions": {},
            "mus": {},
            "sigmas": {},
            "sample_ids": {},
        }
        for mod_name, dynamics in self._modality_dynamics.items():
            sample_ids = batch_data[mod_name].get("sample_ids")
            # get index of
            if sample_ids is not None:
                captured_data["sample_ids"][mod_name] = np.array(sample_ids)

            model_output = dynamics["mp"]
            captured_data["latentspaces"][mod_name] = model_output.latentspace.detach()
            captured_data["reconstructions"][
                mod_name
            ] = model_output.reconstruction.detach()
            if model_output.latent_mean is not None:
                captured_data["mus"][mod_name] = model_output.latent_mean.detach()
            if model_output.latent_logvar is not None:
                captured_data["sigmas"][mod_name] = model_output.latent_logvar.detach()

        return captured_data

    def _dynamics_to_result(self, split: str, epoch_dynamics: List[Dict]) -> None:
        """Aggregates and stores epoch dynamics into the Result object.

        Due to our multi modal Traning with unpaired data, we can see sample_ids more than once per epoch.
        For more consistent downstream analysis, we only keep the first occurence of this sample in the epoch
        and report the dynamics for this sample

        Args:
            split: The data split name (e.g., 'train', 'valid', 'test').
            epoch_dynamics: List of dictionaries capturing dynamics for each batch in the epoch.

        """
        final_data: Dict[str, Any] = defaultdict(lambda: defaultdict(list))
        for batch_data in epoch_dynamics:
            for dynamic_type, mod_data in batch_data.items():
                for mod_name, data in mod_data.items():
                    if isinstance(data, torch.Tensor):
                        data = data.cpu().numpy()
                    final_data[dynamic_type][mod_name].append(data)

        sample_ids: Optional[Dict[str, np.ndarray]] = final_data.get("sample_ids")
        if sample_ids is None:
            raise ValueError("No Sample Ids in TrainingDynamics")
        concat_ids: Dict[str, np.ndarray] = {
            k: np.concatenate(v) for k, v in sample_ids.items()
        }
        unique_idx: Dict[str, np.ndarray] = {
            k: np.unique(v, return_index=True)[1] for k, v in concat_ids.items()
        }

        deduplicated_data: Dict[str, Dict[str, np.ndarray]] = defaultdict(dict)
        # all_dynamics_inner: Dict[str, Any] = {}
        for dynamic_type, dynamic_data in final_data.items():
            concat_dynamic: Dict[str, np.ndarray] = {
                k: np.concatenate(v) for k, v in dynamic_data.items()
            }
            # all_dynamics_inner[dynamic_type] = concat_dynamic
            deduplicated_helper: Dict[str, np.ndarray] = {
                k: v[unique_idx[k]] for k, v in concat_dynamic.items()
            }
            deduplicated_data[dynamic_type] = deduplicated_helper
            # Store the deduplicated data
        # self._all_dynamics[split][self._cur_epoch] = all_dynamics_inner
        self._result.latentspaces.add(
            epoch=self._cur_epoch,
            split=split,
            data=deduplicated_data.get("latentspaces", {}),
        )

        self._result.reconstructions.add(
            epoch=self._cur_epoch,
            split=split,
            data=deduplicated_data.get("reconstructions", {}),
        )

        if deduplicated_data.get("sample_ids"):
            self._result.sample_ids.add(
                epoch=self._cur_epoch,
                split=split,
                data=deduplicated_data["sample_ids"],
            )

        if deduplicated_data.get("mus"):
            self._result.mus.add(
                epoch=self._cur_epoch,
                split=split,
                data=deduplicated_data["mus"],
            )

        if deduplicated_data.get("sigmas"):
            self._result.sigmas.add(
                epoch=self._cur_epoch,
                split=split,
                data=deduplicated_data["sigmas"],
            )

    def _log_losses(
        self,
        split: str,
        total_loss: float,
        sub_losses: Dict[str, float],
        n_samples: int,
    ) -> None:
        """Logs and stores average losses for the epoch.
        Args:
            split: The data split name (e.g., 'train', 'valid').
            total_loss: The cumulative total loss for the epoch.
            sub_losses: Dictionary of accumulated sub-losses for the epoch.
            n_samples: Total number of samples processed in the epoch.

        """

        n_samples = max(n_samples, 1)
        print(f"split: {split}, n_samples: {n_samples}")
        # avg_total_loss = total_loss / n_samples ## Now already normalized
        self._result.losses.add(epoch=self._cur_epoch, split=split, data=total_loss)

        # avg_sub_losses = {
        #     k: v / n_samples if "_factor" not in k else v for k, v in sub_losses.items()
        # }
        self._result.sub_losses.add(
            epoch=self._cur_epoch,
            split=split,
            data=sub_losses,
        )
        self._fabric.print(
            f"Epoch {self._cur_epoch + 1}/{self._config.epochs} - {split.capitalize()} Loss: {total_loss:.4f}"
        )
        # Detailed sub-loss logging in one line
        sub_loss_str = ", ".join(
            [f"{k}: {v:.4f}" for k, v in sub_losses.items() if "_factor" not in k]
        )
        self._fabric.print(f"Sub-losses - {sub_loss_str}")

    def _store_checkpoint(self, split: str, epoch_dynamics: List[Dict]) -> None:
        """Stores model checkpoints and epoch dynamics into the Result object.

        Args:
            split: The data split name (e.g., 'train', 'valid').
            epoch_dynamics: List of dictionaries capturing dynamics for each batch in the epoch.

        """

        state_to_save = {
            mod_name: dynamics["model"].state_dict()
            for mod_name, dynamics in self._modality_dynamics.items()
        }
        state_to_save["latent_clf"] = self._latent_clf.state_dict()
        self._result.model_checkpoints.add(epoch=self._cur_epoch, data=state_to_save)
        self._dynamics_to_result(split=split, epoch_dynamics=epoch_dynamics)

    def _store_final_models(self) -> None:
        """Stores the final trained models into the Result object."""
        final_models = {
            mod_name: dynamics["model"]
            for mod_name, dynamics in self._modality_dynamics.items()
        }
        self._result.model = final_models

    def purge(self) -> None:
        """
        Cleans up all instantiated resources used during training, including
        all modality-specific models/optimizers and the adversarial classifier.
        """

        if hasattr(self, "_modality_dynamics"):
            for mod_name, dynamics in self._modality_dynamics.items():
                if "model" in dynamics and dynamics["model"] is not None:
                    # Remove the model and its optimizer from the dynamics dict
                    del dynamics["model"]
                if "optim" in dynamics and dynamics["optim"] is not None:
                    del dynamics["optim"]
            # Delete the container itself after cleaning contents
            del self._modality_dynamics

        if hasattr(self, "_latent_clf"):
            del self._latent_clf
        if hasattr(self, "_clf_optim"):
            del self._clf_optim

        if hasattr(self, "_trainloader"):
            del self._trainloader
        if hasattr(self, "_validloader") and self._validloader is not None:
            del self._validloader
        if hasattr(self, "_trainset"):
            del self._trainset
        if hasattr(self, "_validset"):
            del self._validset
        if hasattr(self, "_loss_fn"):
            del self._loss_fn
        if hasattr(self, "sub_loss"):  # Directly accessible sub_loss instance
            del self.sub_loss
        if hasattr(self, "_clf_loss_fn"):
            del self._clf_loss_fn

        if hasattr(self, "_all_dynamics"):
            del self._all_dynamics

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        gc.collect()

    def _train_one_epoch_with_profiling(
        self,
    ) -> Tuple[List[Dict], Dict[str, float], int]:
        """Training loop with torch.profiler integration."""
        for dynamics in self._modality_dynamics.values():
            dynamics["model"].train()
        self._latent_clf.train()

        self._clf_epoch_loss = 0
        self._epoch_loss = 0
        epoch_dynamics: List[Dict] = []
        sub_losses: Dict[str, float] = defaultdict(float)
        n_samples_total: int = 0

        # Configure profiler
        with profile(
            activities=[
                ProfilerActivity.CPU,
                ProfilerActivity.CUDA,
            ],
            record_shapes=True,  # Record tensor shapes
            profile_memory=True,  # Track memory usage
            with_stack=True,  # Include stack traces
            # Schedule: skip first 5 batches (warmup), profile next 5, repeat
            schedule=torch.profiler.schedule(
                wait=1,  # Skip first batch (warmup)
                warmup=1,  # Warmup for 1 batch
                active=3,  # Profile 3 batches
                repeat=1,  # Do this once
            ),
            on_trace_ready=torch.profiler.tensorboard_trace_handler(
                "./profiler_logs", worker_name=self._config.profile_logs
            ),
        ) as prof:
            train_iter = iter(self._trainloader)
            for batch_idx, batch in enumerate(self._trainloader):
                with record_function("dataloader"):
                    batch = next(train_iter)
                with record_function("total_batch"):
                    with self._fabric.autocast():
                        # --- Stage 1: Forward for each modality ---
                        with record_function("modalities_forward"):
                            self._modalities_forward(batch=batch)

                        # --- Stage 2: Prepare adversarial training ---
                        with record_function("prep_adver_training"):
                            latents, labels = self._prep_adver_training()
                            n_samples_total += latents.size(0)

                        # --- Stage 3: Train Classifier ---
                        with record_function("train_classifier"):
                            self._train_clf(latents=latents, labels=labels)

                        # --- Stage 4: Train Autoencoders ---
                        with record_function("train_autoencoders"):
                            for _, dynamics in self._modality_dynamics.items():
                                dynamics["optim"].zero_grad()

                            clf_scores_for_adv = self._latent_clf(latents)

                            batch_loss, loss_dict = self._loss_fn(
                                batch=batch,
                                modality_dynamics=self._modality_dynamics,
                                clf_scores=clf_scores_for_adv,
                                labels=labels,
                                clf_loss_fn=self._clf_loss_fn,
                                is_training=True,
                            )

                    with record_function("backward_pass"):
                        self._fabric.backward(batch_loss)

                    with record_function("optimizer_step"):
                        for _, dynamics in self._modality_dynamics.items():
                            dynamics["optim"].step()

                    # --- Logging and Capturing ---
                    self._epoch_loss += batch_loss.item()
                    for k, v in loss_dict.items():
                        value_to_add = v.item() if hasattr(v, "item") else v
                        if "_factor" not in k:
                            sub_losses[k] += value_to_add
                        else:
                            sub_losses[k] = value_to_add

                    if self._is_checkpoint_epoch:
                        with record_function("capture_dynamics"):
                            batch_capture = self._capture_dynamics(batch)
                            epoch_dynamics.append(batch_capture)

                # Step the profiler
                prof.step()

                # Stop profiling after a few batches to avoid huge logs
                if batch_idx >= 5:
                    break
            # Print summary to console
        print("\n" + "=" * 80)
        print("PROFILER SUMMARY - Top Operations by CPU Time")
        print("=" * 80)
        print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=20))

        print("\n" + "=" * 80)
        print("PROFILER SUMMARY - Top Operations by CUDA Time")
        print("=" * 80)
        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

        # --- Write profiler summary to CSV ---

        import csv
        from pathlib import Path

        csv_path = Path(f"./profiler_logs/{self._config.profile_logs}.csv")
        csv_path.parent.mkdir(parents=True, exist_ok=True)

        with csv_path.open(mode="w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(
                [
                    "name",
                    "cpu_timecuda_timecpu_memory_usage",
                    "cuda_memory_usage",
                    "input_shapes",
                    "calls",
                ]
            )

            for evt in prof.key_averages():
                writer.writerow(
                    [
                        evt.key,
                        evt.cpu_time,
                        evt.cuda_time,
                        evt.cpu_memory_usage,
                        evt.cuda_memory_usage,
                        str(evt.input_shapes),
                        evt.count,
                    ]
                )

__init__(trainset, validset, result, config, model_type, loss_type, model_map, sub_loss_type=VarixLoss, ontologies=None, **kwargs)

Initializes the XModalTrainer with datasets and configuration.

Parameters:

Name	Type	Description	Default
trainset	MultiModalDataset	Training dataset containing multiple modalities	required
validset	MultiModalDataset	Validation dataset containing multiple modalities	required
result	Result	Result object to store training outcomes	required
config	DefaultConfig	Configuration parameters for training and model architecture	required
model_type	Type[BaseAutoencoder]	Type of autoencoder model to use for each modality (not used directly)	required
loss_type	Type[BaseLoss]	Type of loss function to use for training the stacked model	required
sub_loss_type	Type[BaseLoss]	Type of loss function to use for individual modality autoencoders	VarixLoss
model_map	Dict[DataSetTypes, Type[BaseAutoencoder]]	Mapping from DataSetTypes to specific autoencoder architectures	required
ontologies	Optional[Union[Tuple, List]]	Ontology information, for compatibility with Ontix	None

Source code in src/autoencodix/trainers/_xmodal_trainer.py

def __init__(
    self,
    trainset: MultiModalDataset,
    validset: MultiModalDataset,
    result: Result,
    config: DefaultConfig,
    model_type: Type[BaseAutoencoder],
    loss_type: Type[BaseLoss],
    model_map: Dict[DataSetTypes, Type[BaseAutoencoder]],
    sub_loss_type: Type[BaseLoss] = VarixLoss,
    ontologies: Optional[Union[Tuple, List]] = None,
    **kwargs,
):
    """Initializes the XModalTrainer with datasets and configuration.

    Args:
        trainset: Training dataset containing multiple modalities
        validset: Validation dataset containing multiple modalities
        result: Result object to store training outcomes
        config: Configuration parameters for training and model architecture
        model_type: Type of autoencoder model to use for each modality (not used directly)
        loss_type: Type of loss function to use for training the stacked model
        sub_loss_type: Type of loss function to use for individual modality autoencoders
        model_map: Mapping from DataSetTypes to specific autoencoder architectures
        ontologies: Ontology information, for compatibility with Ontix
    """
    self._trainset = trainset
    self._n_train_modalities = self._trainset.n_modalities  # type: ignore
    self.model_map = model_map
    self.model_trainer_map = model_trainer_map
    self._n_cpus = os.cpu_count()
    if self._n_cpus is None:
        self._n_cpus = 0

    super().__init__(
        trainset, validset, result, config, model_type, loss_type, ontologies
    )
    # init attributes ------------------------------------------------
    self.latent_dim = config.latent_dim
    self.n_test: Optional[int] = None
    self.n_train = len(trainset.data) if trainset else 0
    self.n_valid = len(validset.data) if validset else 0
    self._cur_epoch: int = 0
    self._is_checkpoint_epoch: Optional[bool] = None
    self._sub_loss_type = sub_loss_type
    self.sub_loss = sub_loss_type(config=self._config)
    self._clf_epoch_loss: float = 0.0
    self._epoch_loss: float = 0.0
    self._epoch_loss_valid: float = 0.0
    self._all_dynamics: Dict[str, Dict[int, Dict]] = {
        "test": defaultdict(dict),
        "train": defaultdict(dict),
        "valid": defaultdict(dict),
    }

decode(x)

Decodes input latent representations Args: x: Latent representations to decode, shape (n_samples, latent_dim)

Source code in src/autoencodix/trainers/_xmodal_trainer.py

def decode(self, x: torch.Tensor):
    """Decodes input latent representations
    Args:
        x: Latent representations to decode, shape (n_samples, latent_dim)
    """
    raise NotImplementedError("General decode step is not implemented for XModalix")

predict(data, model=None, **kwargs)

Performs cross-modal prediction from a specified 'from' modality to a 'to' modality.

The direction is determined by the 'translate_direction' attribute in the config. Results are stored in the Result object under split='test' and epoch=-1.

Parameters:

Name	Type	Description	Default
data	BaseDataset	A MultiModalDataset containing the input data for the 'from' modality.	required

Returns:

Type	Description
Result	The Result object populated with prediction results.

Source code in src/autoencodix/trainers/_xmodal_trainer.py

def predict(
    self,
    data: BaseDataset,
    model: Optional[Dict[str, torch.nn.Module]] = None,
    **kwargs,
) -> Result:
    """Performs cross-modal prediction from a specified 'from' modality to a 'to' modality.

    The direction is determined by the 'translate_direction' attribute in the config.
    Results are stored in the Result object under split='test' and epoch=-1.

    Args:
        data: A MultiModalDataset containing the input data for the 'from' modality.

    Returns:
        The Result object populated with prediction results.
    """
    split: str = kwargs.pop("split", "test")
    if split not in ["train", "valid", "test"]:
        raise ValueError(
            f"split must be one of 'train', 'valid', or 'test', got {split}"
        )
    if not isinstance(data, MultiModalDataset):
        raise TypeError(
            f"type of data has to be MultiModalDataset, got: {type(data)}"
        )
    if model is not None and not hasattr(self, "_modality_dynamics"):
        self._modality_dynamics = {mod_name: {} for mod_name in model.keys()}

        for mod_name, mod_model in model.items():
            mod_model.eval()
            if not isinstance(mod_model, _FabricModule):
                mod_model = self._fabric.setup_module(mod_model)
            mod_model.to(self._fabric.device)
            self._modality_dynamics[mod_name]["model"] = mod_model
        # self._setup_fabric(old_model=model)

    from_key = kwargs.pop("from_key", None)
    to_key = kwargs.pop("to_key", None)
    predict_keys = find_translation_keys(
        config=self._config,
        trained_modalities=list(self._modality_dynamics.keys()),
        from_key=from_key,
        to_key=to_key,
    )
    from_key, to_key = predict_keys["from"], predict_keys["to"]
    from_modality = self._modality_dynamics[from_key]
    to_modality = self._modality_dynamics[to_key]
    from_model = from_modality["model"]
    to_model = to_modality["model"]
    from_model.eval(), to_model.eval()
    # drop_last handled in custom sampler
    inference_loader = DataLoader(
        data,
        batch_sampler=CoverageEnsuringSampler(
            multimodal_dataset=data, batch_size=self._config.batch_size
        ),
        collate_fn=create_multimodal_collate_fn(multimodal_dataset=data),
    )
    inference_loader = self._fabric.setup_dataloaders(inference_loader)  # type: ignore
    epoch_dynamics: List[Dict] = []
    with (
        self._fabric.autocast(),
        torch.inference_mode(),
    ):
        for batch in inference_loader:
            # needed for visualize later
            self._get_vis_dynamics(batch=batch)
            from_z = self._modality_dynamics[from_key]["mp"].latentspace
            translated = to_model.decode(x=from_z)

            # for reference
            to_z = self._modality_dynamics[to_key]["mp"].latentspace
            to_to_reference = to_model.decode(x=to_z)

            batch_capture = self._capture_dynamics(batch)
            translation_key = "translation"

            reference_key = f"reference_{to_key}_to_{to_key}"
            batch_capture["reconstructions"][
                translation_key
            ] = translated.cpu().numpy()

            batch_capture["reconstructions"][
                reference_key
            ] = to_to_reference.cpu().numpy()

            if "sample_ids" in batch[from_key]:
                batch_capture["sample_ids"][translation_key] = np.array(
                    batch[from_key]["sample_ids"]
                )
            if "sample_ids" in batch[to_key]:
                batch_capture["sample_ids"][reference_key] = np.array(
                    batch[to_key]["sample_ids"]
                )

            epoch_dynamics.append(batch_capture)

    self._dynamics_to_result(split=split, epoch_dynamics=epoch_dynamics)

    print("Prediction complete.")
    return self._result

purge()

Cleans up all instantiated resources used during training, including all modality-specific models/optimizers and the adversarial classifier.

Source code in src/autoencodix/trainers/_xmodal_trainer.py

def purge(self) -> None:
    """
    Cleans up all instantiated resources used during training, including
    all modality-specific models/optimizers and the adversarial classifier.
    """

    if hasattr(self, "_modality_dynamics"):
        for mod_name, dynamics in self._modality_dynamics.items():
            if "model" in dynamics and dynamics["model"] is not None:
                # Remove the model and its optimizer from the dynamics dict
                del dynamics["model"]
            if "optim" in dynamics and dynamics["optim"] is not None:
                del dynamics["optim"]
        # Delete the container itself after cleaning contents
        del self._modality_dynamics

    if hasattr(self, "_latent_clf"):
        del self._latent_clf
    if hasattr(self, "_clf_optim"):
        del self._clf_optim

    if hasattr(self, "_trainloader"):
        del self._trainloader
    if hasattr(self, "_validloader") and self._validloader is not None:
        del self._validloader
    if hasattr(self, "_trainset"):
        del self._trainset
    if hasattr(self, "_validset"):
        del self._validset
    if hasattr(self, "_loss_fn"):
        del self._loss_fn
    if hasattr(self, "sub_loss"):  # Directly accessible sub_loss instance
        del self.sub_loss
    if hasattr(self, "_clf_loss_fn"):
        del self._clf_loss_fn

    if hasattr(self, "_all_dynamics"):
        del self._all_dynamics

    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    gc.collect()

train()

Orchestrates the full training process for the cross-modal autoencoder.

Source code in src/autoencodix/trainers/_xmodal_trainer.py

def train(self):
    """Orchestrates the full training process for the cross-modal autoencoder."""
    self._pretraining()
    for epoch in range(self._config.epochs):
        self._cur_epoch = epoch
        self._is_checkpoint_epoch = self._should_checkpoint(epoch=epoch)
        self._fabric.print(f"--- Epoch {epoch + 1}/{self._config.epochs} ---")
        if epoch == 0 and self._config.profiling:
            print("Profiling enabled for epoch 0")
            self._train_one_epoch_with_profiling()
            continue
        train_epoch_dynamics, train_sub_losses, n_samples_train = (
            self._train_one_epoch()
        )

        self._log_losses(
            split="train",
            total_loss=self._epoch_loss,
            sub_losses=train_sub_losses,
            n_samples=n_samples_train,
        )

        if self._validset:
            valid_epoch_dynamics, valid_sub_losses, n_samples_valid = (
                self._validate_one_epoch()
            )
            self._log_losses(
                split="valid",
                total_loss=self._epoch_loss_valid,
                sub_losses=valid_sub_losses,
                n_samples=n_samples_valid,
            )
        if self._is_checkpoint_epoch:
            self._fabric.print(f"Storing checkpoint for epoch {epoch}...")
            self._store_checkpoint(
                split="train", epoch_dynamics=train_epoch_dynamics
            )
            if self._validset:
                self._store_checkpoint(
                    split="valid",
                    epoch_dynamics=valid_epoch_dynamics,
                )
    # save final model
    self._store_checkpoint(split="train", epoch_dynamics=train_epoch_dynamics)
    self._store_final_models()
    return self._result