autoencodix 0.2.0__py3-none-any.whl

autoencodix/utils/__init__.py

@@ -0,0 +1,22 @@
+from ._result import Result
+from ._model_output import ModelOutput
+from ._bulkreader import BulkDataReader
+from ._imgreader import ImageDataReader
+from ._screader import SingleCellDataReader
+from ._losses import VanillixLoss, VarixLoss, DisentanglixLoss, XModalLoss
+from .adata_converter import AnnDataConverter
+
+__all__ = [
+    "Result",
+    "ModelOutput",
+    "BulkDataReader",
+    "ImageDataReader",
+    "SingleCellDataReader",
+    "VanillixLoss",
+    "VarixLoss",
+    "DisentanglixLoss",
+    "XModalLoss",
+    "AnnDataConverter",
+]
+
+# all test done

autoencodix/utils/_annealer.py

@@ -0,0 +1,193 @@
+from typing import Dict, Callable, Optional
+from math import exp
+
+
+class AnnealingScheduler:
+    """VAE loss annealing scheduler with multiple annealing strategies."""
+
+    def __init__(self) -> None:
+        """Initialize the annealing scheduler with strategy mappings."""
+        self._strategies: Dict[str, Callable[[int, int], float]] = {
+            "5phase-constant": self._five_phase_constant,
+            "3phase-linear": self._three_phase_linear,
+            "3phase-log": self._three_phase_log,
+            "logistic-mid": self._logistic_mid,
+            "logistic-early": self._logistic_early,
+            "logistic-late": self._logistic_late,
+            "no-annealing": self._no_annealing,
+        }
+
+    @staticmethod
+    def get_annealing_epoch(
+        *, anneal_pretraining: bool, n_epochs_pretrain: int, current_epoch: int
+    ) -> Optional[int]:
+        """Check if annealing should be used for pretraining.
+        Args:
+            anneal_pretraining: Whether to apply annealing during pretraining phase.
+            n_epochs_pretrain: Number of pretraining epochs.
+            current_epoch: Current epoch number.
+        Returns:
+            int or None: Annealing epoch number, or None if no annealing.
+        Raises:
+            NotImplementedError: This is a deprecated method.
+        """
+        raise NotImplementedError(
+            "Deprecated, for annealing the current epoch is passed, we split between training and \
+                                  pretraining, so no extra calculation is needed"
+        )
+
+    def get_weight(
+        self,
+        *,
+        epoch_current: Optional[int],
+        total_epoch: int,
+        func: str = "logistic-mid",
+    ) -> float:
+        """Calculate VAE loss annealing weight.
+
+        Args:
+            epoch_current: Current epoch in training, or None for full weight.
+            total_epoch: Total epochs for training.
+            func: Specification of annealing function. Default is 'logistic-mid'.
+
+        Returns:
+            Annealing weight between 0 (no VAE loss) and 1 (full VAE loss).
+
+        Raises:
+            NotImplementedError: If the specified annealing function is not implemented.
+        """
+        if epoch_current is None:
+            return 1.0
+
+        if func not in self._strategies:
+            raise NotImplementedError("The annealer is not implemented yet")
+
+        return self._strategies[func](epoch_current, total_epoch)
+
+    def _five_phase_constant(self, epoch_current: int, total_epoch: int) -> float:
+        """Five phase constant annealing strategy.
+
+        Args:
+            epoch_current: Current epoch number.
+            total_epoch: Total number of epochs.
+
+        Returns:
+            Annealing weight.
+        """
+        intervals = 5
+        current_phase = int((epoch_current / total_epoch) * intervals)
+
+        if current_phase == 0:
+            return 0.0
+        elif current_phase == 1:
+            return 0.001
+        elif current_phase == 2:
+            return 0.01
+        elif current_phase == 3:
+            return 0.1
+        else:
+            return 1.0
+
+    def _three_phase_linear(self, epoch_current: int, total_epoch: int) -> float:
+        """Three phase linear annealing strategy.
+
+        Args:
+            epoch_current: Current epoch number.
+            total_epoch: Total number of epochs.
+
+        Returns:
+           Annealing weight.
+        """
+        first_phase_end = total_epoch / 3
+        second_phase_end = 2 * (total_epoch / 3)
+
+        if epoch_current < first_phase_end:
+            return 0.0
+        elif epoch_current < second_phase_end:
+            return (epoch_current - first_phase_end) / (total_epoch / 3)
+        else:
+            return 1.0
+
+    def _three_phase_log(self, epoch_current: int, total_epoch: int) -> float:
+        """Three phase logarithmic annealing strategy.
+
+        Args:
+            epoch_current: Current epoch number.
+            total_epoch: Total number of epochs.
+
+        Returns:
+            Annealing weight.
+        """
+        first_phase_end = total_epoch / 3
+        second_phase_end = 2 * (total_epoch / 3)
+
+        if epoch_current < first_phase_end:
+            return 0.0
+        elif epoch_current < second_phase_end:
+            return self._logistic_mid(epoch_current - first_phase_end, total_epoch / 3)
+        else:
+            return 1.0
+
+    def _compute_logistic_weight(
+        self, epoch_current: int, total_epoch: int, midpoint: float
+    ) -> float:
+        """Compute logistic annealing weight.
+
+        Args:
+            epoch_current: Current epoch number.
+            total_epoch: Total number of epochs.
+            midpoint: Midpoint ratio for the logistic function (0.0 to 1.0).
+
+        Returns:
+            Annealing weight.
+        """
+        b_param = (1 / total_epoch) * 20
+        return 1 / (1 + exp(-b_param * (epoch_current - total_epoch * midpoint)))
+
+    def _logistic_mid(self, epoch_current: int, total_epoch: int) -> float:
+        """Logistic annealing with midpoint at half of total epochs.
+
+        Args:
+            epoch_current: Current epoch number.
+            total_epoch: Total number of epochs.
+
+        Returns:
+            Annealing weight.
+        """
+        return self._compute_logistic_weight(epoch_current, total_epoch, 0.5)
+
+    def _logistic_early(self, epoch_current: int, total_epoch: int) -> float:
+        """Logistic annealing with early midpoint at quarter of total epochs.
+
+        Args:
+            epoch_current: Current epoch number.
+            total_epoch: Total number of epochs.
+
+        Returns:
+            Annealing weight.
+        """
+        return self._compute_logistic_weight(epoch_current, total_epoch, 0.25)
+
+    def _logistic_late(self, epoch_current: int, total_epoch: int) -> float:
+        """Logistic annealing with late midpoint at three quarters of total epochs.
+
+        Args:
+            epoch_current: Current epoch number.
+            total_epoch: Total number of epochs.
+
+        Returns:
+            Annealing weight.
+        """
+        return self._compute_logistic_weight(epoch_current, total_epoch, 0.75)
+
+    def _no_annealing(self, epoch_current: int, total_epoch: int) -> float:
+        """No annealing strategy - constant full weight.
+
+        Args:
+            epoch_current: Current epoch number.
+            total_epoch: Total number of epochs.
+
+        Returns:
+            Annealing weight (always 1.0).
+        """
+        return 1.0

autoencodix/utils/_bulkreader.py

@@ -0,0 +1,232 @@
+import os
+import warnings
+from typing import Dict, Set, Tuple, Optional, Union
+
+import pandas as pd
+
+from autoencodix.configs.default_config import DefaultConfig
+
+
+class BulkDataReader:
+    """Reads bulk data from files based on configuration.
+
+    Supports both paired and unpaired data reading strategies.
+
+    Attributes:
+        config: Configuration object
+    """
+
+    def __init__(self, config: DefaultConfig):
+        """Initialize the BulkDataReader with a configuration.
+
+        Args:
+            config: Configuration object containing data paths and specifications.
+        """
+        self.config = config
+
+    def read_data(self) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
+        """Read all data according to the configuration.
+
+        Returns:
+            A tuple containing (bulk_dataframes, annotation_dataframes)
+        """
+        if self.config.requires_paired or self.config.requires_paired is None:
+            return self.read_paired_data()
+        else:
+            return self.read_unpaired_data()
+
+    def read_paired_data(
+        self,
+    ) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
+        """Reads numeric paired data
+
+        Returns:
+            Tuple containing two Dicts:
+                 1. with name of the data as key and pandas DataFrame as value
+                 2. with  str 'paired' as key and a common annotaion/metadata as DataFrame
+        """
+        common_samples: Optional[Set[str]] = None
+        bulk_dfs: Dict[str, pd.DataFrame] = {}
+        annotation_df = pd.DataFrame()
+        has_annotation = False
+
+        # First pass: read all data files and track common samples
+        for key, info in self.config.data_config.data_info.items():
+            if info.data_type == "IMG":
+                continue
+
+            file_path = os.path.join(info.file_path)
+            df = self._read_tabular_data(file_path, info.sep or "\t")
+
+            if df is None:
+                continue
+
+            if info.data_type == "NUMERIC" and not info.is_single_cell:
+                current_samples = set(df.index)
+                if common_samples is None:
+                    common_samples = current_samples
+                else:
+                    common_samples &= current_samples
+
+                bulk_dfs[key] = df
+
+            elif info.data_type == "ANNOTATION":
+                has_annotation = True
+                annotation_df = df
+
+        # Second pass: filter to common samples
+        if common_samples:
+            common_samples_list = list(common_samples)
+
+            # Reindex bulk dataframes to common samples
+            for key in bulk_dfs:
+                bulk_dfs[key] = bulk_dfs[key].reindex(common_samples_list)
+
+            # Handle annotation dataframe
+            if has_annotation:
+                annotation = annotation_df.reindex(common_samples_list)
+            else:
+                # Create empty annotation with common sample indices
+                annotation_df = pd.DataFrame(index=common_samples_list)
+                annotation = annotation_df
+        else:
+            print("Warning: No common samples found across datasets")
+            annotation = annotation_df
+
+        return bulk_dfs, {"paired": annotation}
+
+    def read_unpaired_data(
+        self,
+    ) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
+        """Read data without enforcing sample alignment across modalities.
+
+        Returns:
+            A tuple containing (bulk_dataframes, annotation_dataframes)
+        """
+        bulk_dfs: Dict[str, pd.DataFrame] = {}
+        annotations: Dict[str, pd.DataFrame] = {}
+
+        for key, info in self.config.data_config.data_info.items():
+            if info.data_type == "IMG" or info.is_single_cell:
+                continue  # Skip image and single-cell data
+
+            # Read main data file
+            file_path = os.path.join(info.file_path)
+            df = self._read_tabular_data(file_path=file_path, sep=info.sep)
+
+            if df is None:
+                continue
+
+            if info.data_type == "NUMERIC":
+                bulk_dfs[key] = df
+
+                if hasattr(info, "extra_anno_file") and info.extra_anno_file:
+                    extra_anno_file = os.path.join(info.extra_anno_file)
+                    extra_anno_df = self._read_tabular_data(
+                        file_path=extra_anno_file, sep=info.sep
+                    )
+                    if extra_anno_df is not None:
+                        annotations[key] = extra_anno_df
+
+            elif info.data_type == "ANNOTATION":
+                annotations[key] = df
+
+        bulk_dfs, annotations = self._validate_and_filter_unpaired(
+            bulk_dfs, annotations
+        )
+
+        return bulk_dfs, annotations
+
+    def _validate_and_filter_unpaired(
+        self,
+        bulk_dfs: Dict[str, pd.DataFrame],
+        annotations: Dict[str, pd.DataFrame],
+    ) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
+        """Validates that all samples in bulk data have a corresponding annotation.
+
+        If a single global annotation file is provided, it creates a perfectly
+        matched annotation dataframe for each bulk dataframe.
+
+        Warns and drops samples that do not have a corresponding annotation.
+
+        Args:
+            bulk_dfs: Dictionary of bulk data modalities and their dataframes.
+            annotations: Dictionary of annotation dataframes, possibly one global one.
+
+        Returns:
+            A tuple of two dictionaries:
+            1. The filtered bulk dataframes.
+            2. The new, synchronized annotation dataframes, with keys matching the bulk dataframes.
+        """
+        if not annotations:
+            warnings.warn(
+                "No annotation files were provided. Cannot validate sample annotations."
+            )
+            return bulk_dfs, {}
+
+        # If annotations have keys that match bulk_dfs, we assume they are already paired.
+        # This logic focuses on the case where one annotation file is meant for all bulk files.
+        # A simple heuristic: if there is one annotation file and its key is not in bulk_dfs.
+        annotation_keys = set(annotations.keys())
+        bulk_keys = set(bulk_dfs.keys())
+
+        # Check for the global annotation case
+        if len(annotation_keys) == 1 and not annotation_keys.intersection(bulk_keys):
+            global_annotation_key = list(annotation_keys)[0]
+            global_annotation_df = annotations[global_annotation_key]
+
+            filtered_bulk_dfs = {}
+            synchronized_annotations = {}
+
+            for key, data_df in bulk_dfs.items():
+                data_samples = data_df.index
+                annotation_samples = global_annotation_df.index
+
+                # Find the intersection of valid sample IDs
+                valid_ids = data_samples.intersection(annotation_samples)
+
+                # Check for and warn about dropped samples
+                if len(valid_ids) < len(data_samples):
+                    missing_ids = sorted(list(set(data_samples) - set(valid_ids)))
+                    warnings.warn(
+                        f"For data modality '{key}', {len(missing_ids)} sample(s) "
+                        f"were found without a corresponding annotation and will be dropped: {missing_ids}"
+                    )
+
+                # Filter both the data and the annotation to the valid IDs
+                filtered_bulk_dfs[key] = data_df.loc[valid_ids]
+                synchronized_annotations[key] = global_annotation_df.loc[valid_ids]
+
+            return filtered_bulk_dfs, synchronized_annotations
+        else:
+            # Handle the case where annotations are already meant to be paired by key
+            # (Or a more complex case we are not handling yet)
+            warnings.warn(
+                "Proceeding without global annotation synchronization. Assuming annotations are pre-aligned by key."
+            )
+            return bulk_dfs, annotations
+
+    def _read_tabular_data(
+        self, file_path: str, sep: Union[str, None] = None
+    ) -> pd.DataFrame:
+        """Read tabular data from a file with error handling.
+
+        Args:
+        file_path: Path to the data file.
+        sep: Separator character for CSV/TSV files.
+
+        Returns:
+            The loaded DataFrame.
+        """
+        try:
+            if file_path.endswith(".parquet"):
+                print(f"reading parquet: {file_path}")
+                return pd.read_parquet(file_path)
+            elif file_path.endswith((".csv", ".txt", ".tsv")):
+                return pd.read_csv(file_path, sep=sep, index_col=0)
+            else:
+                raise ValueError(
+                    f"Unsupported file type for {file_path}. Supported formats: .parquet, .csv, .txt, .tsv"
+                )
+        except Exception as e:
+            raise e

autoencodix/utils/_explainer.py

@@ -0,0 +1,246 @@
+import torch
+import numpy as np
+import scipy
+import pandas as pd
+from autoencodix.modeling._captum_forward import CaptumForward
+from typing import Optional, Union
+from captum.attr import (
+    DeepLiftShap,
+    IntegratedGradients,
+)
+import warnings
+
+warnings.filterwarnings(
+    "ignore",
+    message="Setting forward, backward hooks and attributes on non-linear",
+    category=UserWarning,
+)
+
+warnings.filterwarnings("ignore")
+
+
+class FeatureImportanceExplainer:
+    """
+    More complete version with:
+      - method selection (DeepLiftShap, IG, etc.)
+      - baseline construction (mean / random / grouped)
+      - subset sampling
+      - reproducible randomness
+    """
+
+    def __init__(
+        self,
+        adata,
+        model,
+        method: str = "DeepLiftShap",
+        sel_latent_dim: Union[
+            list, int, None
+        ] = None,  # list or int of latent dimensions to explain, if None explain all
+        n_subset: int = 100,  # Randomly sample in input and baseline space for computational efficiency, if n_subset is None use all samples
+        seed_int: int = 12,
+        input_type: str = "random",  # "random" or "grouped" for selecting subset of inputs to explain
+        input_group: Optional[
+            str
+        ] = None,  # column in .obs for grouping inputs must be provided if input_type is "grouped"
+        baseline_type: str = "random",  # "random" or "mean" or "grouped" for selecting baseline samples
+        baseline_group: str = None,  # column in .obs for grouping inputs must be provided if input_type is "grouped", optional for "mean", ignored for "random"
+        anno_col: Optional[str] = None,  # column in .obs for grouping
+    ):
+        super(FeatureImportanceExplainer, self).__init__()
+        self.adata_ACX = adata
+        self.model = model
+        self.latent_dim = model.config.latent_dim
+
+        self.method = method
+        if self.method not in {"DeepLiftShap", "IntegratedGradients"}:
+            raise ValueError(f"Invalid method {method}.")
+        self.n_subset = n_subset
+        self.seed_int = seed_int
+        self.baseline_type = baseline_type
+        self.baseline_group = baseline_group
+        self.anno_col = anno_col
+        self.input_type = input_type
+        self.input_group = input_group
+        self.sel_latent_dim = sel_latent_dim
+
+        torch.manual_seed(seed_int)
+        np.random.seed(seed_int)
+
+    def explain(self):
+        adata_ACX = self.adata_ACX
+        gene_names = adata_ACX.var_names
+        inputs, baselines = return_inputs_baseline(
+            adata_ACX,
+            self.baseline_group,
+            self.baseline_type,
+            self.anno_col,
+            self.input_type,
+            self.input_group,
+        )
+        # set n_samples to whatever is lowest from: inputs, baselines, n_subset (if not None)
+        n_samples = min(inputs.shape[0], baselines.shape[0])
+        if self.n_subset is not None:
+            n_samples = min(n_samples, self.n_subset)
+
+        # Downsample input
+        if self.n_subset is not None:
+            indices_keep = np.random.choice(
+                inputs.shape[0], size=n_samples, replace=False
+            )
+        else:
+            indices_keep = np.arange(inputs.shape[0])
+        # Downsample baselines in the same way
+        if self.n_subset is not None:
+            indices_keep_baseline = np.random.choice(
+                baselines.shape[0], size=n_samples, replace=False
+            )
+        else:
+            indices_keep_baseline = np.arange(baselines.shape[0])
+        # for latent_dim in range(self.latent_dim):
+        all_attr = []
+        if isinstance(
+            self.sel_latent_dim, list
+        ):  # Provide list of latent dimensions to explain
+            latent_dims = self.sel_latent_dim
+        elif isinstance(
+            self.sel_latent_dim, int
+        ):  # Provide single latent dimension to explain
+            latent_dims = [self.sel_latent_dim]
+        elif self.sel_latent_dim is None:  # Explain all latent dimensions
+            latent_dims = list(range(self.latent_dim))
+        else:
+            raise ValueError(
+                f"Invalid sel_latent_dim {self.sel_latent_dim}. Must be int, list of ints, or None."
+            )
+
+        for latent_dim in latent_dims:
+            print("Calculating attributions for latent dimension:", latent_dim)
+            cp_forward_dim = CaptumForward(model=self.model, dim=latent_dim)
+            if self.method == "DeepLiftShap":
+                cp_explainer = DeepLiftShap(cp_forward_dim)
+            if self.method == "IntegratedGradients":
+                cp_explainer = IntegratedGradients(cp_forward_dim)
+            attributions = cp_explainer.attribute(
+                inputs=inputs[indices_keep].float(),
+                baselines=baselines[indices_keep_baseline].float(),
+                return_convergence_delta=False,
+            )
+            avg_abs_attributions = attributions.abs().mean(dim=0)
+            all_attr.append(avg_abs_attributions.detach().cpu())
+        attr_matrix = torch.stack(all_attr).T.numpy()
+        if hasattr(self.model, "ontologies") and self.model.ontologies is not None:
+            cols = [list(self.model.ontologies[0].keys())[i] for i in latent_dims]
+        else:
+            cols = [f"Latent_Dim_{i}" for i in latent_dims]
+        df_attributions = pd.DataFrame(
+            attr_matrix,
+            index=list(gene_names),
+            columns=cols,
+        )
+
+        return df_attributions
+
+
+def return_inputs_baseline(
+    adata, baseline_group, baseline_type, anno_col, input_type, input_group
+):
+    baseline_torch_all = torch.tensor(
+        adata.X.toarray() if scipy.sparse.issparse(adata.X) else adata.X
+    )
+    ## Define inputs
+    if input_type == "grouped":
+        if anno_col is None or input_group is None:
+            raise ValueError(
+                "If input_type is 'grouped', anno_col and input_group must be provided to specify the grouping column in .obs."
+            )
+        input_adata = adata[adata.obs[anno_col] == input_group]
+        inputs = torch.tensor(
+            input_adata.X.toarray()
+            if scipy.sparse.issparse(input_adata.X)
+            else input_adata.X
+        )
+    elif input_type == "random":
+        inputs = torch.tensor(
+            adata.X.toarray() if scipy.sparse.issparse(adata.X) else adata.X
+        )
+
+    ## Define baselines
+    if baseline_type == "random":
+        # Generates a bootstrap random from input in the same size
+        baseline_random = baseline_torch_all[
+            torch.randint(0, baseline_torch_all.size(0), (1,)).item()
+        ]
+        baselines = torch.tensor(
+            np.tile(baseline_random, (baseline_torch_all.shape[0], 1))
+        )
+    elif baseline_type == "grouped":
+        if anno_col is None or baseline_group is None:
+            raise ValueError(
+                "If baseline_type is 'grouped', anno_col and baseline_group must be provided to specify the grouping column in .obs."
+            )
+        base_adata_filtered = adata[adata.obs[anno_col] == baseline_group]
+        base_filtered = torch.tensor(
+            base_adata_filtered.X.toarray()
+            if scipy.sparse.issparse(base_adata_filtered.X)
+            else base_adata_filtered.X
+        )
+        # Generates a bootstrap random from the filtered input in the same size
+        baseline_grouped = base_filtered[
+            torch.randint(0, base_filtered.size(0), (1,)).item()
+        ]
+        baselines = torch.tensor(
+            np.tile(baseline_grouped, (base_adata_filtered.shape[0], 1))
+        )
+    elif baseline_type == "mean":
+        if baseline_group == None:
+            baseline_mean = baseline_torch_all.mean(axis=0)  # gene_means
+        else:
+            if anno_col is None:
+                raise ValueError(
+                    "If baseline_group is not 'all', anno_col must be provided to specify the grouping column in .obs."
+                )
+            base_adata_filtered = adata[adata.obs[anno_col] == baseline_group]
+            base_filtered = torch.tensor(
+                base_adata_filtered.X.toarray()
+                if scipy.sparse.issparse(base_adata_filtered.X)
+                else base_adata_filtered.X
+            )
+            baseline_mean = base_filtered.mean(axis=0)  # gene_means
+        baselines = torch.tensor(
+            np.tile(baseline_mean, (baseline_torch_all.shape[0], 1))
+        )
+    else:
+        raise ValueError(f"Invalid baseline_type {baseline_type}.")
+
+    # inputs = torch.tensor(
+    #     input_adata.X.toarray()
+    #     if scipy.sparse.issparse(input_adata.X)
+    #     else input_adata.X
+    # )
+    # if baseline_group == "all":
+    #     if baseline_type == "mean":
+    #         baseline_mean = inputs.mean(axis=0)  # gene_means
+    #         baselines = torch.tensor(np.tile(baseline_mean, (inputs.shape[0], 1)))
+    #     if baseline_type == "random_sample":
+    #         baseline_random = inputs[torch.randint(0, inputs.size(0), (1,)).item()]
+    #         baselines = torch.tensor(np.tile(baseline_random, (inputs.shape[0], 1)))
+    # else:
+    #     if anno_col is None:
+    #         raise ValueError(
+    #             "If baseline_group is not 'all', anno_col must be provided to specify the grouping column in .obs."
+    #         )
+    #     input_adata_filtered = input_adata[input_adata.obs[anno_col] == baseline_group]
+    #     inputs_filtered = torch.tensor(
+    #         input_adata_filtered.X.toarray()
+    #         if scipy.sparse.issparse(input_adata_filtered.X)
+    #         else input_adata_filtered.X
+    #     )
+    #     if baseline_type == "mean":
+    #         baseline_mean = inputs_filtered.mean(axis=0)  # gene_means
+    #         baselines = torch.tensor(np.tile(baseline_mean, (inputs.shape[0], 1)))
+    #     if baseline_type == "random_sample":
+    #         baseline_random = inputs_filtered[
+    #             torch.randint(0, inputs_filtered.size(0), (1,)).item()
+    #         ]
+    #         baselines = torch.tensor(np.tile(baseline_random, (inputs.shape[0], 1)))
+    return inputs, baselines