fusion-bench 0.2.20__py3-none-any.whl → 0.2.21__py3-none-any.whl

fusion_bench/__init__.py

@@ -20,6 +20,7 @@ from . import (
     utils,
 )
 from .method import BaseAlgorithm, BaseModelFusionAlgorithm
+from .mixins import auto_register_config
 from .modelpool import BaseModelPool
 from .models import separate_io
 from .taskpool import BaseTaskPool

fusion_bench/_get_started/__init__.py

@@ -0,0 +1,3 @@
+"""
+Tutorial module for FusionBench
+"""

fusion_bench/_get_started/greeting_program.py

@@ -0,0 +1,49 @@
+import logging
+from typing import Optional
+
+from omegaconf import DictConfig
+
+from fusion_bench.programs import BaseHydraProgram
+
+log = logging.getLogger(__name__)
+
+
+class GreetingProgram(BaseHydraProgram):
+    """
+    A simple program that greets users with a custom message.
+    """
+
+    _config_mapping = BaseHydraProgram._config_mapping | {
+        "message": "message",
+        "name": "name",
+        "repeat_count": "repeat_count",
+    }
+
+    def __init__(
+        self,
+        message: str = "Hello",
+        name: str = "World",
+        repeat_count: int = 1,
+        **kwargs,
+    ):
+        self.message = message
+        self.name = name
+        self.repeat_count = repeat_count
+        super().__init__(**kwargs)
+
+    def run(self):
+        """Execute the greeting workflow."""
+        log.info("Starting greeting program")
+
+        # Create the greeting
+        greeting = f"{self.message}, {self.name}!"
+
+        # Print the greeting multiple times
+        for i in range(self.repeat_count):
+            if self.repeat_count > 1:
+                print(f"[{i+1}/{self.repeat_count}] {greeting}")
+            else:
+                print(greeting)
+
+        log.info("Greeting program completed")
+        return greeting

fusion_bench/compat/method/base_algorithm.py

@@ -36,6 +36,20 @@ class ModelFusionAlgorithm(ABC):
             algorithm_config = DictConfig({})
         self.config = algorithm_config
 
+    def on_run_start(self):
+        """
+        Hook method called at the start of the run.
+        Can be overridden by subclasses to perform initialization tasks.
+        """
+        pass
+
+    def on_run_end(self):
+        """
+        Hook method called at the end of the run.
+        Can be overridden by subclasses to perform cleanup tasks.
+        """
+        pass
+
     @abstractmethod
     def run(self, modelpool):
         """

fusion_bench/constants/__init__.py

@@ -1,2 +1,7 @@
 # flake8: noqa F401
+import importlib.metadata
+
 from .paths import *
+
+# fusionbench version
+FUSION_BENCH_VERSION = importlib.metadata.version("fusion-bench")

fusion_bench/constants/clip_vision.py

@@ -1,4 +1,5 @@
-# Constants for CLIP Vision Model Merging
+"Constants for CLIP Vision Model Merging"
+
 TASK_NAMES_TA8 = [
     "sun397",
     "stanford-cars",
@@ -9,7 +10,23 @@ TASK_NAMES_TA8 = [
     "mnist",
     "dtd",
 ]
-
+"The 8 tasks used in the Task Arithmetic paper."
+TASK_NAMES_TALL8 = TASK_NAMES_TA8
+"The 8 tasks used in the Tall Mask paper"
+TASK_NAMES_TALL10 = TASK_NAMES_TA8 + ["oxford_flowers102", "pcam"]
+TASK_NAMES_TALL12 = TASK_NAMES_TALL10 + [
+    "fer2013",
+    "oxford-iiit-pet",
+]
+TASK_NAMES_TALL14 = TASK_NAMES_TALL12 + [
+    "stl10",
+    "cifar100",
+]
+"The 14 tasks used in the TALL mask paper"
+TASK_NAMES_TALL16 = TASK_NAMES_TALL14 + ["cifar10", "food101"]
+TASK_NAMES_TALL18 = TASK_NAMES_TALL16 + ["fashion_mnist", "emnist_letters"]
+TASK_NAMES_TALL20 = TASK_NAMES_TALL18 + ["kmnist", "rendered-sst2"]
+"The 20 tasks used in the TALL mask paper"
 TASK_NAMES_TA8_CAP = [
     "SUN397",
     "Cars",
@@ -20,3 +37,10 @@ TASK_NAMES_TA8_CAP = [
     "MNIST",
     "DTD",
 ]
+TASK_NAMES_TALL8_CAP = TASK_NAMES_TA8_CAP
+TASK_NAMES_TALL10_CAP = TASK_NAMES_TALL8_CAP + ["Flowers102", "PCAM"]
+TASK_NAMES_TALL12_CAP = TASK_NAMES_TALL10_CAP + ["FER2013", "OxfordIIITPet"]
+TASK_NAMES_TALL14_CAP = TASK_NAMES_TALL12_CAP + ["STL10", "CIFAR100"]
+TASK_NAMES_TALL16_CAP = TASK_NAMES_TALL14_CAP + ["CIFAR10", "Food101"]
+TASK_NAMES_TALL18_CAP = TASK_NAMES_TALL16_CAP + ["FashionMNIST", "EMNIST"]
+TASK_NAMES_TALL20_CAP = TASK_NAMES_TALL18_CAP + ["KMNIST", "RenderedSST2"]

fusion_bench/constants/paths.py

@@ -7,10 +7,14 @@ log = logging.getLogger(__name__)
 __all__ = ["LIBRARY_PATH", "PROJECT_ROOT_PATH", "DEFAULT_CONFIG_PATH"]
 
 LIBRARY_PATH = Path(importlib.import_module("fusion_bench").__path__[0])
+"""Path to the library directory."""
+
 PROJECT_ROOT_PATH = LIBRARY_PATH.parent
+"""Path to the project root directory."""
 
 if (PROJECT_ROOT_PATH / "config").is_dir():
     DEFAULT_CONFIG_PATH = PROJECT_ROOT_PATH / "config"
+    """Path to the default config directory."""
 elif (PROJECT_ROOT_PATH / "fusion_bench_config").is_dir():
     DEFAULT_CONFIG_PATH = PROJECT_ROOT_PATH / "fusion_bench_config"
 else:

fusion_bench/dataset/clip_dataset.py

@@ -5,6 +5,7 @@ This module provides a class to convert a dataset whose object is a list of dict
 from typing import Optional, Tuple
 
 import torch
+from torch.utils.data import Dataset
 from transformers import CLIPProcessor, ProcessorMixin
 
 __all__ = ["CLIPDataset"]
@@ -28,7 +29,7 @@ class CLIPDataset(torch.utils.data.Dataset):
         processor (CLIPProcessor): The CLIP processor used for image preprocessing.
     """
 
-    def __init__(self, dataset, processor: Optional[CLIPProcessor] = None):
+    def __init__(self, dataset: Dataset, processor: Optional[CLIPProcessor] = None):
         self.dataset = dataset
         self.processor = processor
 

fusion_bench/dataset/gpt2_glue.py

@@ -16,7 +16,7 @@ from functools import partial
 from pathlib import Path
 from typing import Literal
 
-from datasets import load_dataset, load_from_disk
+from datasets import Dataset, load_dataset, load_from_disk
 from transformers import PreTrainedTokenizer
 
 
@@ -147,7 +147,7 @@ class TokenizedGLUE:
         return glue_dataset_loaders[name]()
 
     @cache_dataset
-    def load_mrpc_dataset(self):
+    def load_mrpc_dataset(self) -> Dataset:
         """
         Load and tokenize the MRPC dataset.
 
@@ -166,7 +166,7 @@ class TokenizedGLUE:
         return dataset
 
     @cache_dataset
-    def load_rte_dataset(self):
+    def load_rte_dataset(self) -> Dataset:
         """
         Load and tokenize the RTE dataset.
 
@@ -186,7 +186,7 @@ class TokenizedGLUE:
         return dataset
 
     @cache_dataset
-    def load_wnli_dataset(self):
+    def load_wnli_dataset(self) -> Dataset:
         """
         Load and tokenize the WNLI dataset.
 
@@ -205,7 +205,7 @@ class TokenizedGLUE:
         return dataset
 
     @cache_dataset
-    def load_qqp_dataset(self):
+    def load_qqp_dataset(self) -> Dataset:
         """
         Load and tokenize the QQP dataset.
 
@@ -224,7 +224,7 @@ class TokenizedGLUE:
         return dataset
 
     @cache_dataset
-    def load_mnli_dataset(self):
+    def load_mnli_dataset(self) -> Dataset:
         """
         Load and tokenize the MNLI dataset.
 
@@ -243,7 +243,7 @@ class TokenizedGLUE:
         return dataset
 
     @cache_dataset
-    def load_cola_dataset(self):
+    def load_cola_dataset(self) -> Dataset:
         """
         Load and tokenize the CoLA dataset.
 
@@ -262,7 +262,7 @@ class TokenizedGLUE:
         return dataset
 
     @cache_dataset
-    def load_sst2_dataset(self):
+    def load_sst2_dataset(self) -> Dataset:
         """
         Load and tokenize the SST-2 dataset.
 
@@ -281,7 +281,7 @@ class TokenizedGLUE:
         return dataset
 
     @cache_dataset
-    def load_qnli_dataset(self):
+    def load_qnli_dataset(self) -> Dataset:
         """
         Load and tokenize the QNLI dataset.
 

fusion_bench/dataset/image_corruption/make_corruption.py

@@ -0,0 +1,179 @@
+# -*- coding: utf-8 -*-
+import logging
+
+logger = logging.getLogger(__name__)
+
+import collections
+import warnings
+from io import BytesIO
+
+import cv2  # pip install opencv-python
+import numpy as np
+import skimage as sk
+import torch
+import torchvision.transforms as trn
+from PIL import Image
+from PIL import Image as PILImage
+from scipy.ndimage import zoom as scizoom
+from scipy.ndimage.interpolation import map_coordinates
+from skimage.filters import gaussian  # pip install scikit-image
+from tqdm import tqdm
+
+try:
+    from wand.api import library as wandlibrary
+    from wand.image import Image as WandImage
+except ImportError as e:
+    logger.error(
+        "Failed to import wand."
+        "Install it with `apt-get install libmagickwand-dev` and `pip install Wand`"
+        "For more information, refer to the documentation https://docs.wand-py.org/"
+    )
+    raise e
+
+# /////////////// Distortion Helpers ///////////////
+
+warnings.simplefilter("ignore", UserWarning)
+
+
+# /////////////// Distortions ///////////////
+class MotionImage(WandImage):
+    def motion_blur(self, radius=0.0, sigma=0.0, angle=0.0):
+        wandlibrary.MagickMotionBlurImage(self.wand, radius, sigma, angle)
+
+
+def gaussian_noise(x, severity=1):
+    c = [0.04, 0.06, 0.08, 0.09, 0.10][severity - 1]
+
+    x = np.array(x) / 255.0
+    return np.clip(x + np.random.normal(size=x.shape, scale=c), 0, 1) * 255
+
+
+def impulse_noise(x, severity=1):
+    c = [0.01, 0.02, 0.03, 0.05, 0.07][severity - 1]
+
+    x = sk.util.random_noise(np.array(x) / 255.0, mode="s&p", amount=c)
+    return np.clip(x, 0, 1) * 255
+
+
+def motion_blur(x, severity=1):
+    c = [(6, 1), (6, 1.5), (6, 2), (8, 2), (9, 2.5)][severity - 1]
+
+    output = BytesIO()
+    x.save(output, format="PNG")
+    x = MotionImage(blob=output.getvalue())
+
+    x.motion_blur(radius=c[0], sigma=c[1], angle=np.random.uniform(-45, 45))
+
+    x = cv2.imdecode(np.fromstring(x.make_blob(), np.uint8), cv2.IMREAD_UNCHANGED)
+
+    if x.shape != (32, 32):
+        return np.clip(x[..., [2, 1, 0]], 0, 255)  # BGR to RGB
+    else:  # greyscale to RGB
+        return np.clip(np.array([x, x, x]).transpose((1, 2, 0)), 0, 255)
+
+
+def spatter(x, severity=1):
+    c = [
+        (0.62, 0.1, 0.7, 0.7, 0.5, 0),
+        (0.65, 0.1, 0.8, 0.7, 0.5, 0),
+        (0.65, 0.3, 1, 0.69, 0.5, 0),
+        (0.65, 0.1, 0.7, 0.69, 0.6, 1),
+        (0.65, 0.1, 0.5, 0.68, 0.6, 1),
+    ][severity - 1]
+    x = np.array(x, dtype=np.float32) / 255.0
+
+    liquid_layer = np.random.normal(size=x.shape[:2], loc=c[0], scale=c[1])
+
+    liquid_layer = gaussian(liquid_layer, sigma=c[2])
+    liquid_layer[liquid_layer < c[3]] = 0
+    if c[5] == 0:
+        liquid_layer = (liquid_layer * 255).astype(np.uint8)
+        dist = 255 - cv2.Canny(liquid_layer, 50, 150)
+        dist = cv2.distanceTransform(dist, cv2.DIST_L2, 5)
+        _, dist = cv2.threshold(dist, 20, 20, cv2.THRESH_TRUNC)
+        dist = cv2.blur(dist, (3, 3)).astype(np.uint8)
+        dist = cv2.equalizeHist(dist)
+        #     ker = np.array([[-1,-2,-3],[-2,0,0],[-3,0,1]], dtype=np.float32)
+        #     ker -= np.mean(ker)
+        ker = np.array([[-2, -1, 0], [-1, 1, 1], [0, 1, 2]])
+        dist = cv2.filter2D(dist, cv2.CV_8U, ker)
+        dist = cv2.blur(dist, (3, 3)).astype(np.float32)
+
+        m = cv2.cvtColor(liquid_layer * dist, cv2.COLOR_GRAY2BGRA)
+        m /= np.max(m, axis=(0, 1))
+        m *= c[4]
+
+        # water is pale turqouise
+        color = np.concatenate(
+            (
+                175 / 255.0 * np.ones_like(m[..., :1]),
+                238 / 255.0 * np.ones_like(m[..., :1]),
+                238 / 255.0 * np.ones_like(m[..., :1]),
+            ),
+            axis=2,
+        )
+
+        color = cv2.cvtColor(color, cv2.COLOR_BGR2BGRA)
+        x = cv2.cvtColor(x, cv2.COLOR_BGR2BGRA)
+
+        return cv2.cvtColor(np.clip(x + m * color, 0, 1), cv2.COLOR_BGRA2BGR) * 255
+    else:
+        m = np.where(liquid_layer > c[3], 1, 0)
+        m = gaussian(m.astype(np.float32), sigma=c[4])
+        m[m < 0.8] = 0
+        #         m = np.abs(m) ** (1/c[4])
+
+        # mud brown
+        color = np.concatenate(
+            (
+                63 / 255.0 * np.ones_like(x[..., :1]),
+                42 / 255.0 * np.ones_like(x[..., :1]),
+                20 / 255.0 * np.ones_like(x[..., :1]),
+            ),
+            axis=2,
+        )
+
+        color *= m[..., np.newaxis]
+        x *= 1 - m[..., np.newaxis]
+
+        return np.clip(x + color, 0, 1) * 255
+
+
+def contrast(x, severity=1):
+    c = [0.75, 0.5, 0.4, 0.3, 0.15][severity - 1]
+
+    x = np.array(x) / 255.0
+    means = np.mean(x, axis=(0, 1), keepdims=True)
+    return np.clip((x - means) * c + means, 0, 1) * 255
+
+
+def jpeg_compression(x, severity=1):
+    c = [80, 65, 58, 50, 40][severity - 1]
+
+    output = BytesIO()
+    x.save(output, "JPEG", quality=c)
+    x = PILImage.open(output)
+
+    return x
+
+
+def pixelate(x, severity=1):
+    c = [0.95, 0.9, 0.85, 0.75, 0.65][severity - 1]
+
+    x = x.resize((int(32 * c), int(32 * c)), PILImage.BOX)
+    x = x.resize((32, 32), PILImage.BOX)
+
+    return x
+
+
+# /////////////// End Distortions ///////////////
+
+
+distortion_methods = collections.OrderedDict()
+distortion_methods["Gaussian Noise"] = gaussian_noise
+distortion_methods["Impulse Noise"] = impulse_noise
+distortion_methods["Motion Blur"] = motion_blur
+distortion_methods["Contrast"] = contrast
+distortion_methods["Pixelate"] = pixelate
+distortion_methods["JPEG"] = jpeg_compression
+distortion_methods["Spatter"] = spatter

fusion_bench/dataset/image_dataset.py

@@ -20,7 +20,7 @@ class TransformedImageDataset(Dataset):
         transform (Callable): The transform to be applied to the images.
     """
 
-    def __init__(self, dataset, transform: Callable):
+    def __init__(self, dataset: Dataset, transform: Callable):
         super().__init__()
         self.dataset = dataset
         self.transform = transform

fusion_bench/dataset/nyuv2.py

@@ -1,6 +1,6 @@
 import fnmatch
 import os
-from typing import Callable, Optional
+from typing import Callable, Dict, Optional, Tuple
 
 import numpy as np
 import torch
@@ -68,7 +68,7 @@ class NYUv2(Dataset):
         )
         self.noise = torch.rand(self.data_len, 1, 288, 384)
 
-    def __getitem__(self, index):
+    def __getitem__(self, index) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
         """
         Retrieve an item from the dataset.
 

fusion_bench/method/__init__.py

@@ -37,11 +37,12 @@ _import_structure = {
     "ties_merging": ["TiesMergingAlgorithm"],
     "dare": ["DareSimpleAverage", "DareTaskArithmetic", "DareTiesMerging"],
     "fisher_merging": [
+        "FisherMergingAlgorithm",
         "FisherMergingForCLIPVisionModel",
         "FisherMergingAlgorithmForGPT2",
     ],
     "regmean": ["RegMeanAlgorithmForCLIP", "RegMeanAlgorithmForGPT2"],
-    "regmean_plusplus": ["RegMeanAlgorithmForCLIPPlusPlus"],
+    "regmean_plusplus": ["RegMeanAlgorithmPlusPlus", "RegMeanAlgorithmForCLIPPlusPlus"],
     "adamerging": [
         "CLIPTaskWiseAdaMergingAlgorithm",
         "CLIPLayerWiseAdaMergingAlgorithm",
@@ -69,6 +70,7 @@ _import_structure = {
         "FlanT5LayerWiseGossipAlgorithm",
     ],
     "fw_merging": ["FrankWolfeHardAlgorithm", "FrankWolfeSoftAlgorithm"],
+    "tall_mask": ["TallMaskTaskArithmeticAlgorithm"],
     # plug-and-play model merging methods
     "concrete_subspace": [
         "ConcreteTaskArithmeticAlgorithmForCLIP",
@@ -99,6 +101,8 @@ _import_structure = {
         "SmileUpscalingAlgorithm",
         "SingularProjectionMergingAlgorithm",
     ],
+    # task vector compression methods
+    "bitdelta": ["BitDeltaAlgorithm"],
     # pruning methods
     "pruning": [
         "MagnitudeDiffPruningAlgorithm",
@@ -126,6 +130,7 @@ if TYPE_CHECKING:
     from .adamerging import *
     from .analysis import TaskVectorCosSimilarity, TaskVectorViolinPlot
     from .base_algorithm import BaseAlgorithm, BaseModelFusionAlgorithm
+    from .bitdelta import BitDeltaAlgorithm
     from .classification import (
         ContinualImageClassificationFineTuningForCLIP,
         ImageClassificationFineTuningForCLIP,
@@ -154,7 +159,11 @@ if TYPE_CHECKING:
         LayerWisePruningForMixtral,
         ProgressivePruningForMixtral,
     )
-    from .fisher_merging import FisherMergingForCLIPVisionModel
+    from .fisher_merging import (
+        FisherMergingAlgorithm,
+        FisherMergingAlgorithmForGPT2,
+        FisherMergingForCLIPVisionModel,
+    )
     from .fw_merging import FrankWolfeHardAlgorithm, FrankWolfeSoftAlgorithm
     from .gossip import (
         CLIPLayerWiseGossipAlgorithm,
@@ -196,7 +205,10 @@ if TYPE_CHECKING:
     )
     from .rankone_moe import CLIPRankOneMoEAlgorithm, RankOneMoEAlgorithm
     from .regmean import RegMeanAlgorithmForCLIP, RegMeanAlgorithmForGPT2
-    from .regmean_plusplus import RegMeanAlgorithmForCLIPPlusPlus
+    from .regmean_plusplus import (
+        RegMeanAlgorithmForCLIPPlusPlus,
+        RegMeanAlgorithmPlusPlus,
+    )
     from .simple_average import SimpleAverageAlgorithm
     from .slerp import SlerpMergeAlgorithm
     from .smile_upscaling import (
@@ -212,6 +224,7 @@ if TYPE_CHECKING:
         PCPSparseLoForLlama,
         SparseLoForLlama,
     )
+    from .tall_mask import TallMaskTaskArithmeticAlgorithm
     from .task_arithmetic import TaskArithmeticAlgorithm
     from .task_singular_vector import TaskSingularVectorMerging
     from .ties_merging import TiesMergingAlgorithm

fusion_bench/method/adamerging/clip_layer_wise_adamerging.py

@@ -3,7 +3,7 @@ Example Usage:
 
 ```bash
 fusion_bench \
-    method=adamerging \
+    method=adamerging/clip \
         method.name=clip_layer_wise_adamerging \
         method.save_merging_weights=merging_weights.pt \
     modelpool=clip-vit-base-patch32_TA8 \

fusion_bench/method/adamerging/clip_task_wise_adamerging.py

@@ -1,6 +1,7 @@
 import functools
 import logging
 import os
+from typing import Iterator
 
 import torch
 from omegaconf import DictConfig
@@ -42,7 +43,7 @@ class CLIPTaskWiseAdaMergingAlgorithm(TaskWiseAdaMergingAlgorithm):
         super().__init__(algorithm_config)
 
     @functools.cache
-    def get_test_dataset(self, task: str):
+    def get_test_dataset(self, task: str) -> CLIPDataset:
         """
         Load the test dataset for the task.
         This method is cached, so the dataset is loaded only once.
@@ -59,7 +60,7 @@ class CLIPTaskWiseAdaMergingAlgorithm(TaskWiseAdaMergingAlgorithm):
         return dataset
 
     @functools.cache
-    def get_shuffled_test_loader_iter(self, task: str):
+    def get_shuffled_test_loader_iter(self, task: str) -> Iterator:
         """
         Get an iterator over the shuffled test DataLoader for the task.
 
@@ -88,11 +89,14 @@ class CLIPTaskWiseAdaMergingAlgorithm(TaskWiseAdaMergingAlgorithm):
         classification head for each task.
         """
         clip_model_config = self.modelpool.get_model_config("_pretrained_")
-        pretrained_path = (
-            clip_model_config.pretrained_model_name_or_path
-            if hasattr(clip_model_config, "pretrained_model_name_or_path")
-            else clip_model_config.path
-        )
+        if isinstance(clip_model_config, str):
+            pretrained_path = clip_model_config
+        else:
+            pretrained_path = (
+                clip_model_config.pretrained_model_name_or_path
+                if hasattr(clip_model_config, "pretrained_model_name_or_path")
+                else clip_model_config.path
+            )
 
         with timeit_context("Loading CLIP processor and pretrained CLIP model."):
             self._clip_processor = CLIPProcessor.from_pretrained(pretrained_path)

fusion_bench/method/adamerging/layer_wise_adamerging.py

@@ -31,9 +31,9 @@ log = logging.getLogger(__name__)
 
 
 class LayerWiseAdaMergingAlgorithm(
-    ModelFusionAlgorithm,
     LightningFabricMixin,
     SimpleProfilerMixin,
+    ModelFusionAlgorithm,
 ):
     _program: "FabricModelFusionProgram"
     """The program that this algorithm is running on."""
@@ -55,7 +55,9 @@ class LayerWiseAdaMergingAlgorithm(
         super().__init__(algorithm_config)
 
     @torch.no_grad()
-    def construct_layer_wise_merged_model(self, modelpool: "ModelPool"):
+    def construct_layer_wise_merged_model(
+        self, modelpool: "ModelPool"
+    ) -> LayerWiseMergedModel:
         """
         Constructs a wrapped layer-wise merged model from model pool.
 
@@ -125,7 +127,7 @@ class LayerWiseAdaMergingAlgorithm(
                 os.makedirs(os.path.dirname(save_path), exist_ok=True)
             torch.save(merging_weights.detach().cpu(), save_path)
 
-    def run(self, modelpool: ModelPool, **kwargs):
+    def run(self, modelpool: ModelPool, **kwargs) -> nn.Module:
         """
         Run the Layer-Wise AdaMerging Algorithm.
 
@@ -176,7 +178,9 @@ class LayerWiseAdaMergingAlgorithm(
         pass
 
     @abstractmethod
-    def compute_logits(self, module, images: Tensor, task: str) -> Tensor:
+    def compute_logits(
+        self, module: LayerWiseMergedModel, images: Tensor, task: str
+    ) -> Tensor:
         """
         Compute the logits for the given images and task.
 
@@ -190,7 +194,9 @@ class LayerWiseAdaMergingAlgorithm(
         """
         pass
 
-    def test_time_adaptation(self, module: "LayerWiseMergedModel[TorchModelType]"):
+    def test_time_adaptation(
+        self, module: "LayerWiseMergedModel[TorchModelType]"
+    ) -> "LayerWiseMergedModel[TorchModelType]":
         """
         Perform test-time adaptation on the merged model.