Documentation for the model

def add_organism(
    self, organism: str, genes: pd.Index, emb: pd.DataFrame, locs=None
):
    """
    Add a new organism to an existing model for transfer learning.

    Extends the gene vocabulary and embeddings to include genes from a new
    organism. Useful for applying a pretrained model to a new species.

    Args:
        organism (str): Organism ontology term ID (e.g., "NCBITaxon:10090" for mouse).
        genes (pd.Index): Gene names/IDs for the new organism.
        emb (pd.DataFrame): Gene embeddings DataFrame with genes as index.
            Will be resized to match model's d_model.
        locs (pd.DataFrame, optional): Genomic positions with 'pos' column.
            Required if model uses positional encoding. Defaults to None.

    Raises:
        ValueError: If model requires gene locations but none provided.
        ValueError: If gene positions exceed model's maximum position encoding.

    Note:
        Only genes present in both `genes` and `emb` (and `locs` if provided)
        will be added. The model's gene encoder is expanded in-place.
    """
    if self.pos_encoder is not None and locs is None:
        raise ValueError("this model needs gene locations to add a new organism")

    self.organisms.append(organism)
    if locs is not None:
        overlap = set(locs.index) & set(emb.index) & set(genes.index)
        genes = genes[genes.index.isin(overlap)]
        locs = locs.loc[genes.index]
        pos = locs["pos"]
        token_to_pos = {token: pos for token, pos in enumerate(pos)}
        if self.pos_encoder.pe.shape[0] < max(pos):
            raise ValueError(
                f"the number of gene locs in the added organism needs to be less than {self.pos_encoder.pe.shape[0]}"
            )

        token_to_pos = {token: pos for token, pos in enumerate(pos)}
        arr = []
        for _, v in token_to_pos.items():
            arr.append(self.pos_encoder.pe[v - 1].to("cpu").numpy())
        pe = torch.Tensor(np.array(arr)).to(self.pos_encoder.pe.device)
        self.pos_encoder.pe = torch.cat([self.pos_encoder.pe, pe], dim=0)
    else:
        overlap = set(emb.index) & set(genes.index)
        genes = genes[genes.index.isin(overlap)]

    emb = emb.loc[genes.index]
    self._genes[organism] = genes.index.tolist()
    if self.gene_encoder is None:
        genc = self.expr_encoder.gene_encoder
    else:
        genc = self.gene_encoder
    if type(genc) is torch.nn.Sequential:
        enc = genc[0]
    else:
        enc = genc
    semb = torch.nn.AdaptiveAvgPool1d(self.d_model)(
        torch.tensor(emb.values, dtype=torch.float32)
    ).to(enc.embeddings.weight.data.device)
    if enc.memmap:
        print("todev.. will fail for now")

    embs = torch.cat([enc.embeddings.weight.data, semb], dim=0)
    enc.embeddings = nn.Embedding(
        embs.shape[0],
        embs.shape[1],
        padding_idx=None,
        _freeze=enc.embeddings.weight.requires_grad,
    )
    enc.embeddings.weight.data.copy_(embs)
    enc.embeddings.weight.data = enc.embeddings.weight.data.to(self.device)
    if type(genc) is torch.nn.Sequential:
        genc[0] = enc
    else:
        genc = enc
    if self.gene_encoder is None:
        self.expr_encoder.gene_encoder = genc
    else:
        self.gene_encoder = genc

def configure_optimizers(self):
    """@see pl.LightningModule"""
    # https://pytorch.org/docs/stable/generated/torch.optim.Adam.html#torch.optim.Adam
    # not working because of poor weight decay implem
    if self.optim == "adam":
        optimizer = optim.Adam(
            self.parameters(),
            lr=self.hparams.lr,
            betas=(0.9, 0.999),
            eps=1e-7,  # 1e-5 to 1e-8
            weight_decay=self.weight_decay,
            amsgrad=False,
            fused=self.fused_adam,
        )
    elif self.optim == "adamW":
        optimizer = optim.AdamW(
            self.parameters(),
            lr=self.hparams.lr,
            betas=(0.9, 0.999),
            eps=1e-7,  # 1e-5 to 1e-8
            weight_decay=self.weight_decay,
            amsgrad=False,
            fused=self.fused_adam,
        )
    elif self.optim == "galore":
        raise NotImplementedError("Galore optimizer not implemented")
        # param_groups = [
        #    {
        #        "params": [
        #            v for k, v in self.named_parameters() if "transformer" not in k
        #        ]
        #    },
        #    {
        #        "params": [
        #            v for k, v in self.named_parameters() if "transformer" in k
        #        ],
        #        "rank": 128,
        #        "update_proj_gap": 200,
        #        "scale": 0.25,
        #        "proj_type": "std",
        #    },
        # ]
        # optimizer = GaLoreAdamW(param_groups, lr=self.hparams.lr)
    else:
        raise ValueError(f"Unknown optimizer: {self.optim}")
    if self.lr_reduce_monitor is None:
        print("no lr reduce factor")
        return [optimizer]
    # lr_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
    #    optimizer,
    #    T_0=20000,
    #    T_mult=2,
    #    eta_min=1e-8,
    # )
    # interval = "step"
    # frequency = 10
    # lr_scheduler = optim.lr_scheduler.ExponentialLR(
    #    optimizer,
    #    gamma=0.85,
    # )
    lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode="min",
        patience=self.lr_reduce_patience,
        factor=self.lr_reduce_factor,
    )
    interval = "epoch"
    frequency = 1
    # lr_scheduler = StepwiseCAWRWithWD(
    #     optimizer,
    #     T_0=20_000,
    #     T_mult=2,
    #     eta_min=1e-8,
    #     wd_decay=0.9
    # )
    lr_dict = {
        "scheduler": lr_scheduler,
        # The unit of the scheduler's step size, could also be 'step'.
        # 'epoch' updates the scheduler on epoch end whereas 'step'
        # updates it after a optimizer update.
        "interval": interval,
        # How many epochs/steps should pass between calls to
        # `scheduler.step()`. 1 corresponds to updating the learning
        # rate after every epoch/step.
        "frequency": frequency,
        # Metric to to monitor for schedulers like `ReduceLROnPlateau`
        "monitor": self.lr_reduce_monitor,
    }
    self.lrfinder_steps = 0
    for val in self.trainer.callbacks:
        if type(val) is _LRCallback:
            self.lrfinder_steps = val.num_training
        if type(val) is LearningRateFinder:
            self.lrfinder_steps = val._num_training_steps
    return [optimizer], [lr_dict]

def forward(
    self,
    gene_pos: Tensor,
    expression: Optional[Tensor] = None,
    neighbors: Optional[Tensor] = None,
    neighbors_info: Optional[Tensor] = None,
    mask: Optional[Tensor] = None,
    req_depth: Optional[Tensor] = None,
    get_gene_emb: bool = False,
    metacell_token: Optional[Tensor] = None,  # (minibatch, 1)
    depth_mult: Optional[Tensor] = None,
    do_sample: bool = False,
    do_mvc: bool = False,
    do_class: bool = False,
    get_attention_layer: Optional[list] = None,
    mask_zeros: Optional[Tensor] = None,
) -> Dict[str, Tensor] | tuple[Dict[str, Tensor], list]:
    """
    Complete forward pass through the scPRINT-2 model.

    Encodes input expression data, processes through transformer(s), and
    decodes into expression predictions and cell classifications.

    Args:
        gene_pos (Tensor): Gene indices of shape (batch, seq_len) mapping to
            positions in the model's gene vocabulary.
        expression (Tensor, optional): Expression values of shape (batch, seq_len).
            Can be raw counts or normalized depending on model config.
            Defaults to None.
        neighbors (Tensor, optional): KNN neighbor expressions of shape
            (batch, n_neighbors, seq_len) for metacell-style encoding.
            Defaults to None.
        neighbors_info (Tensor, optional): Neighbor weights of shape
            (batch, n_neighbors). Defaults to None.
        mask (Tensor, optional): Boolean mask of shape (batch, seq_len) where
            True indicates positions to mask (set to zero). Defaults to None.
        req_depth (Tensor, optional): Target sequencing depth of shape (batch,)
            for depth-conditional generation. Defaults to None.
        get_gene_emb (bool, optional): Return gene embeddings from transformer.
            Defaults to False.
        metacell_token (Tensor, optional): Binary metacell indicators of shape
            (batch,). Defaults to None.
        depth_mult (Tensor, optional): Expression depth multiplier. If None,
            uses sum of expression values. Defaults to None.
        do_sample (bool, optional): Sample from predicted distribution.
            Currently unused. Defaults to False.
        do_mvc (bool, optional): Compute multi-view coding predictions.
            Defaults to False.
        do_class (bool, optional): Compute classification predictions.
            Defaults to False.
        get_attention_layer (list[int], optional): Layer indices to extract
            attention weights from. Defaults to None.
        mask_zeros (Tensor, optional): Boolean mask for zero-expression genes
            of shape (batch, seq_len + num_special_tokens). Defaults to None.

    Returns:
        dict[str, Tensor] | tuple[dict, list]: Model outputs containing:
            - "mean": Predicted expression (batch, seq_len)
            - "disp": Dispersion parameters (batch, seq_len) [if ZINB]
            - "zero_logits": Zero-inflation logits (batch, seq_len) [if ZINB]
            - "input_cell_embs": Cell embeddings (batch, n_classes+1, d_model)
            - "input_cell_emb": Mean cell embedding (batch, d_model)
            - "output_cell_embs": Processed cell embeddings
            - "output_cell_emb": Final cell embedding
            - "cls_output_{class}": Classification logits for each class
            - "gene_embedding": Gene embeddings [if get_gene_emb]
            - "mvc_*": MVC predictions [if do_mvc]

            If get_attention_layer is not None, returns (outputs_dict, attention_list)
            where attention_list contains QKV tensors from specified layers.

    Example:
        >>> output = model(
        ...     gene_pos=batch["genes"],
        ...     expression=batch["x"],
        ...     req_depth=batch["depth"],
        ...     do_class=True,
        ... )
        >>> predictions = output["mean"]
        >>> cell_types = output["cls_output_cell_type_ontology_term_id"].argmax(-1)
    """
    cell_embs, encoding = self._encoder(
        gene_pos,
        expression,
        neighbors,
        neighbors_info,
        mask,
        metacell_token=metacell_token,
    )

    # attention_bias
    num = (1 if self.use_metacell_token else 0) + (
        (len(self.classes) + 1) if not self.cell_transformer else 0
    )
    if self.attn_bias is not None:
        if not hasattr(self, "nbias_sparse"):
            bias_path = os.path.join(self.attn_bias)
            # Keep as sparse matrix - much more memory efficient
            self.nbias_sparse = load_npz(bias_path)

        bias = torch.zeros(
            (
                gene_pos.shape[0],
                gene_pos.shape[1] + num,
                gene_pos.shape[1] + num,
            ),
            device=gene_pos.device,
            dtype=torch.float16,
        )

        fade_factor = 100

        # Extract only the needed values from sparse matrix
        batch_size = gene_pos.shape[0]

        # Vectorized extraction from sparse matrix
        for b in range(batch_size):
            indices = gene_pos[b].cpu().numpy()
            # Get submatrix for this batch's genes
            submatrix = self.nbias_sparse[np.ix_(indices, indices)]
            bias[b, num:, num:] = (
                torch.tensor(
                    submatrix.toarray(), device=gene_pos.device, dtype=torch.float16
                )
                * fade_factor
            )

        bias[:, num:, :num] = -10_000
    if not self.cell_transformer:
        encoding = torch.cat([cell_embs, encoding], dim=1)
    if type(self.transformer) is FlashTransformer:
        transformer_output = self.transformer(
            encoding,
            return_qkv=get_attention_layer,
            bias=bias if self.attn_bias is not None else None,
            bias_layer=list(range(self.nlayers - 1)),
            mask_zeros=mask_zeros,
        )
    elif type(self.transformer) is Performer:
        transformer_output = self.transformer(encoding)
    else:
        raise ValueError(f"Unknown transformer: {type(self.transformer)}")
    if get_attention_layer is not None:
        transformer_output, qkvs = transformer_output
    if self.cell_transformer:
        cell_embs = self.cell_transformer(cell_embs, x_kv=transformer_output)
    else:
        cell_embs, transformer_output = transformer_output.split(
            [
                len(self.classes) + 1,
                transformer_output.shape[1] - (len(self.classes) + 1),
            ],
            dim=1,
        )
    # if not provided we will mult by the current expression sum
    depth_mult = expression.sum(1) if depth_mult is None else depth_mult
    req_depth = torch.log2(1 + req_depth)
    res = self._expr_decoder(
        transformer_output[:, (1 if self.use_metacell_token else 0) :, :],
        depth_mult,
        req_depth,
        get_gene_emb,
    )
    res.update(
        self._cell_decoder(
            cell_embs,
            do_mvc,
            do_class,
            depth_mult,
            req_depth,
            gene_pos if do_mvc else None,
        )
    )
    return (res, qkvs) if get_attention_layer is not None else res

def log_adata(self, gtclass=None, name=""):
    """
    log_adata will log an adata from predictions.
    It will log to tensorboard and wandb if available

    see @utils.log_adata
    """
    try:
        mdir = self.logger.save_dir if self.logger.save_dir is not None else "/tmp"
    except:
        mdir = "data/"
    if not os.path.exists(mdir):
        os.makedirs(mdir)
    adata, fig = utils.make_adata(
        genes=self.genes,
        embs=self.embs,
        pos=self.pos if self.save_expr else None,
        expr_pred=self.expr_pred if self.save_expr else None,
        classes=self.classes,
        pred=self.pred if not self.keep_all_labels_pred else None,
        label_decoders=self.label_decoders,
        labels_hierarchy=self.labels_hierarchy,
        gtclass=gtclass,
        doplot=self.doplot,
    )
    adata.write(
        str(mdir)
        + "/step_"
        + str(self.global_step)
        + "_"
        + str(self.name)
        + "_"
        + str(name)
        + "_"
        + str(self.global_rank)
        + ".h5ad"
    )
    if self.doplot and fig is not None:
        logged = False
        try:
            self.logger.experiment.add_figure(fig)
            logged = True
        except:
            print("couldn't log to tensorboard")
        try:
            self.logger.log_image(key="umaps", images=[fig], step=self.global_step)
            logged = True
        except:
            print("couldn't log to wandb")
        if not logged:
            fig.savefig(mdir + "/umap_" + self.name + "_" + name + ".png")

    return adata

def on_fit_start(self):
    """@see pl.LightningModule"""
    if type(self.transformer) is FlashTransformer:
        for encoder_layers in self.transformer.blocks:
            encoder_layers.set_seq_parallel(True)
    for k, v in self.mat_labels_hierarchy.items():
        self.mat_labels_hierarchy[k] = v.to(self.device)

def on_load_checkpoint(self, checkpoints):
    """
    Handle checkpoint loading with backward compatibility.

    Automatically handles:
    - Different class configurations between checkpoint and current model
    - Legacy parameter names and structures
    - Encoder/decoder mismatches with datamodule
    - Gene vocabulary differences
    - Early stopping callback state

    Called automatically by PyTorch Lightning during checkpoint loading.

    Args:
        checkpoints (dict): Checkpoint dictionary from torch.load().

    Note:
        Prints warnings when configurations differ between checkpoint and
        current model. These should be reviewed to ensure expected behavior.
    """
    # if not the same number of labels (due to diff datasets)
    for name, clss in self.cls_decoders.items():
        size = checkpoints["state_dict"][
            "cls_decoders." + name + ".out_layer.bias"
        ].shape[0]
        if size != clss.out_layer.bias.shape[0]:
            self.cls_decoders[name].out_layer = torch.nn.Linear(
                clss.out_layer.weight.shape[1], size
            )
    # from older model versions
    self.normalization = checkpoints["hyper_parameters"].get("normalization", "sum")
    if (
        checkpoints["state_dict"].get("gene_encoder.0.embedding.weight", None)
        is not None
    ):
        # replace it with the new one gene_encoder.0.embeddings.weight in the state_dict
        checkpoints["state_dict"]["gene_encoder.0.embeddings.weight"] = checkpoints[
            "state_dict"
        ]["gene_encoder.0.embedding.weight"]
        del checkpoints["state_dict"]["gene_encoder.0.embedding.weight"]
    # same
    # when doing batch effect correction and input dataset is not the same
    if (
        "grad_reverse_discriminator_loss.out_layer.bias"
        in checkpoints["state_dict"]
    ):
        for k in list(checkpoints["state_dict"].keys()):
            if "grad_reverse_discriminator_loss" in k:
                del checkpoints["state_dict"][k]
        print(
            "the discriminator for batch effect correction has been removed. "
            "dropping the legacy key."
        )
    # same
    if (
        checkpoints["state_dict"].get("gene_encoder.embedding.weight", None)
        is not None
    ):
        # replace it with the new one gene_encoder.embeddings.weight in the state_dict
        checkpoints["state_dict"]["gene_encoder.embeddings.weight"] = checkpoints[
            "state_dict"
        ]["gene_encoder.embedding.weight"]
        del checkpoints["state_dict"]["gene_encoder.embedding.weight"]

    if "classes" in checkpoints["hyper_parameters"]:
        if self.label_counts != checkpoints["hyper_parameters"]["classes"]:
            if "label_counts" in checkpoints["hyper_parameters"] and set(
                checkpoints["hyper_parameters"]["label_counts"].keys()
            ) == set(checkpoints["hyper_parameters"]["classes"]):
                if self.classes != checkpoints["hyper_parameters"]["classes"]:
                    print("classes have changed, be careful")
                self.classes = checkpoints["hyper_parameters"]["classes"]
                self.label_counts = checkpoints["hyper_parameters"]["label_counts"]
                if self.classes == self.label_counts:
                    raise ValueError(
                        "classes and label_counts are the same, this is not allowed, please use another checkpoint"
                    )
            else:
                self.label_counts = checkpoints["hyper_parameters"]["classes"]
                if self.classes != list(
                    checkpoints["hyper_parameters"]["classes"].keys()
                ):
                    print("classes have changed, be careful")
                    self.classes = list(
                        checkpoints["hyper_parameters"]["classes"].keys()
                    )
        # else it is all good as expected

    else:
        print("no classes in the checkpoint, be careful")

    if checkpoints["state_dict"].get("pos_encoder.pe") is not None:
        if self.pos_encoder is None:
            self.pos_encoder = encoders.PositionalEncoding(
                self.d_model, gene_pos_enc=[0, 1, 2]
            )
        self.pos_encoder.pe = checkpoints["state_dict"]["pos_encoder.pe"]

    if self.label_decoders != checkpoints["hyper_parameters"][
        "label_decoders"
    ] or self.labels_hierarchy != checkpoints["hyper_parameters"].get(
        "labels_hierarchy", {}
    ):
        print("label decoders have changed, be careful")
        self.label_decoders = checkpoints["hyper_parameters"]["label_decoders"]
        self.labels_hierarchy = checkpoints["hyper_parameters"].get(
            "labels_hierarchy", {}
        )
        for k, v in self.labels_hierarchy.items():
            tens = torch.zeros((len(v), self.label_counts[k]))
            for k2, v2 in v.items():
                tens[k2 - self.label_counts[k], v2] = 1
            self.mat_labels_hierarchy[k] = tens.to(bool)

    if (
        "gene_pos_enc" in checkpoints["hyper_parameters"]
        and checkpoints["hyper_parameters"]["gene_pos_enc"] is not None
    ):
        if (
            self.pos_encoder is None
            or self.pos_encoder.gene_pos_enc
            != checkpoints["hyper_parameters"]["gene_pos_enc"]
        ):
            print(
                "Gene position encoding has changed in the dataloader compared to last time, trying to revert"
            )
            self.pos_encoder = encoders.PositionalEncoding(
                self.d_model,
                gene_pos_enc=checkpoints["hyper_parameters"]["gene_pos_enc"],
            )
            checkpoints["hyper_parameters"].pop("gene_pos_enc")
    mencoders = {}
    if type(checkpoints["hyper_parameters"]["genes"]) is list:
        print("converting a gene list-based model")
        org = checkpoints["hyper_parameters"].get("organisms", self.organisms)
        genedf = load_genes(org)
        checkpoints["hyper_parameters"]["genes"] = {
            i: genedf.index[
                (genedf.organism == i)
                & genedf.index.isin(checkpoints["hyper_parameters"]["genes"])
            ].tolist()
            for i in org
        }
    if "precpt_gene_emb" in checkpoints["hyper_parameters"]:
        checkpoints["hyper_parameters"].pop("precpt_gene_emb")

    if "gene_pos_file" in checkpoints["hyper_parameters"]:
        checkpoints["hyper_parameters"].pop("gene_pos_file")

    if "transformer" in checkpoints["hyper_parameters"]:
        checkpoints["hyper_parameters"]["attention"] = checkpoints[
            "hyper_parameters"
        ].pop("transformer")
    try:
        if self.trainer.datamodule.decoders != self.label_decoders:
            print("label decoders have changed, be careful")
            # if we don't have the same decoders, we need to update the one on the datamodule side
            for k, v in self.label_decoders.items():
                mencoders[k] = {va: ke for ke, va in v.items()}
            self.trainer.datamodule.encoders = mencoders

        es = None
        for k in self.trainer.callbacks:
            if isinstance(k, EarlyStopping):
                es = k
        if es is not None:
            prev = checkpoints["callbacks"].get(
                "EarlyStopping{'monitor': 'val_loss', 'mode': 'min'}"
            )
            if prev is not None:
                prev = prev["patience"]
            if prev != es.patience:
                print(
                    "updating the early stopping parameter to {}".format(
                        es.patience
                    )
                )
                checkpoints["callbacks"][
                    "EarlyStopping{'monitor': 'val_loss', 'mode': 'min'}"
                ]["patience"] = es.patience
                if prev < es.patience:
                    checkpoints["callbacks"][
                        "EarlyStopping{'monitor': 'val_loss', 'mode': 'min'}"
                    ]["stopped_epoch"] = 0

    except RuntimeError as e:
        if "scPRINT2 is not attached to a `Trainer`." in str(e):
            print("FYI: scPRINT2 is not attached to a `Trainer`.")
        else:
            raise e
    if (
        self.mvc_decoder is None
        and checkpoints["state_dict"].get("mvc_decoder.gene2query.weight")
        is not None
    ):
        for i in [
            "mvc_decoder.gene2query.weight",
            "mvc_decoder.gene2query.bias",
            "mvc_decoder.norm.weight",
            "mvc_decoder.norm.bias",
            "mvc_decoder.pred_var_zero.weight",
        ]:
            if i in checkpoints["state_dict"]:
                del checkpoints["state_dict"][i]
    org = checkpoints["hyper_parameters"].get("organisms")
    if self.organisms != org and org is not None:
        self.organisms = org
        try:
            self.trainer.datamodule.organisms = self.organisms
        except RuntimeError as e:
            if "scPRINT2 is not attached to a `Trainer`." not in str(e):
                raise e
    if self._genes != checkpoints["hyper_parameters"]["genes"]:
        self._genes = checkpoints["hyper_parameters"]["genes"]
    try:
        self.trainer.datamodule.set_valid_genes_collator(self.genes)
    except RuntimeError as e:
        if "scPRINT2 is not attached to a `Trainer`." not in str(e):
            raise e

    if not is_interactive():
        self.save_hyperparameters()

def on_predict_epoch_end(self):
    """@see pl.LightningModule will"""
    if self.pos.shape[0] < 100:
        return
    if self.pred_log_adata:
        print("adding on disk")
        return self.log_adata(name="predict_part_" + str(self.counter))

def on_predict_epoch_start(self):
    """@see pl.LightningModule"""
    print("predict epoch start")
    self.embs = None
    self.attn.data = None
    self.attn.attn = None
    self.counter = 0
    if type(self.transformer) is FlashTransformer:
        for encoder_layers in self.transformer.blocks:
            encoder_layers.set_seq_parallel(False)

def on_test_start(self):
    """@see pl.LightningModule"""
    print("test start")
    for k, v in self.mat_labels_hierarchy.items():
        self.mat_labels_hierarchy[k] = v.to(self.device)

def on_validation_epoch_end(self):
    """@see pl.LightningModule"""
    self.pos = None
    self.expr_pred = None
    self.do_adv_cls = self._store_adv_cls
    gathered_embs = self.all_gather(self.embs)
    # Merge the dictionaries from all processes
    for key in self.embs.keys():
        self.embs[key] = gathered_embs[key].view(-1, gathered_embs[key].shape[-1])
    self.info = self.all_gather(self.info).view(-1, self.info.shape[-1])
    self.pred = (
        self.all_gather(self.pred).view(-1, self.pred.shape[-1])
        if self.pred is not None
        else None
    )
    # self.pos = self.all_gather(self.pos).view(-1, self.pos.shape[-1])
    # self.expr_pred[0] = self.all_gather(self.expr_pred[0]).view(
    #     -1, self.expr_pred[0].shape[-1]
    # )
    # if len(self.expr_pred) > 1:
    #     self.expr_pred[1] = self.all_gather(self.expr_pred[1]).view(
    #         -1, self.expr_pred[1].shape[-1]
    #     )
    # self.expr_pred[2] = self.all_gather(self.expr_pred[2]).view(
    #     -1, self.expr_pred[2].shape[-1]
    # )

    if self.trainer.state.stage != "sanity_check":
        if self.trainer.is_global_zero:
            print("logging anndata")
            sch = self.lr_schedulers()
            if sch is not None:
                sch.step(self.trainer.callback_metrics["val_loss"])
            # run the test function on specific dataset
            if self.embs is not None:
                self.log_adata(
                    gtclass=self.info, name="validation_part_" + str(self.counter)
                )
            if (self.current_epoch + 1) % self.test_every == 0:
                self.on_test_epoch_end()
            # Synchronize all processes with a timeout
        if torch.distributed.is_initialized():
            # Set a timeout that's longer than your test typically takes
            # Write rank to file for debugging
            self.trainer.strategy.barrier()
    self.pred = None