model commit

README.md

@@ -1,3 +1,43 @@
----
-license: cc-by-nc-4.0
----
+---
+license: cc-by-nc-4.0
+---
+
+# TAP-CT: 3D Task-Agnostic Pretraining of CT Foundation Models
+
+TAP-CT is a suite of foundation models for computed tomography (CT) imaging, pretrained in a task-agnostic manner through an adaptation of DINOv2 for volumetric data. These models learn robust 3D representations from CT scans without requiring task-specific annotations. 
+
+This repository provides TAP-CT-S-2D, a Vision Transformer (ViT-Small) architecture pretrained on volumetric inputs with a spatial resolution of (224, 224) and a patch size of (16, 16). For inference on full-resolution CT volumes, a sliding window approach can be employed to extract features across the entire scan. Additional TAP-CT model variants, as well as the image processor, will be released in future updates.
+
+## Preprocessing
+
+While a dedicated image processor will be released in future updates, optimal feature extraction requires the following preprocessing pipeline:
+
+1. **Orientation**: Convert the volume to LPS (Left-Posterior-Superior) orientation. While the model is likely orientation-invariant, all evaluations were conducted using LPS orientation.
+2. **Spatial Resizing**: Resize the volume to a spatial resolution of \(z, 224, 224\) or \(z, 512, 512\), where \(z\) represents the number of slices along the axial dimension.
+3. **Intensity Clipping**: Clip voxel intensities to the range \([-1008, 822]\) HU (Hounsfield Units).
+4. **Normalization**: Apply z-score normalization using \(mean = -86.8086\) and \(std = 322.6347\).
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModel
+
+# Load the model
+model = AutoModel.from_pretrained('fomofo/tap-ct-s-2d', trust_remote_code=True)
+
+# Prepare input (batch_size, channels, height, width)
+x = torch.randn((16, 1, 224, 224))
+
+# Forward pass
+output = model.forward(x)
+```
+
+The model returns a `BaseModelOutputWithPooling` object from the transformers library. The `output.pooler_output` contains the pooled `[CLS]` token representation, while `output.last_hidden_state` contains the spatial patch token embeddings. To extract features from all intermediate transformer layers, pass `output_hidden_states=True` to the forward method.
+
+## Model Details
+
+- **Model Type**: 3D CT Vision Foundation Model
+- **Input Shape**: `(batch_size, 1, height, width)`
+- **Example Input**: `(16, 1, 224, 224)` - batch of 16 CT slices at 224×224 resolution
+- **License**: CC-BY-NC-4.0

attention.py

@@ -0,0 +1,201 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+import logging
+import os
+import warnings
+from typing import Optional
+
+import torch
+from torch import nn
+
+logger = logging.getLogger("dinov2")
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import memory_efficient_attention, unbind
+
+        XFORMERS_AVAILABLE = True
+        warnings.warn("xFormers is available (Attention)")
+    else:
+        warnings.warn("xFormers is disabled (Attention)")
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+    warnings.warn("xFormers is not available (Attention)")
+
+
+class Attention(nn.Module):
+    """Multi-head self-attention module.
+
+    Parameters
+    ----------
+    dim : int
+        Dimension of the input features.
+    num_heads : int, optional
+        Number of attention heads, by default 8.
+    qkv_bias : bool, optional
+        Whether to add a bias to the query, key, and value projections, by default False.
+    proj_bias : bool, optional
+        Whether to add a bias to the output projection, by default True.
+    attn_drop : float, optional
+        Dropout rate for the attention weights, by default 0.0.
+    proj_drop : float, optional
+        Dropout rate for the output projection, by default 0.0.
+
+    Raises
+    ------
+    ValueError
+        If `dim` is not divisible by `num_heads`.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        """Inits :class:`Attention`.
+
+        Parameters
+        ----------
+        dim : int
+            Dimension of the input features.
+        num_heads : int, optional
+            Number of attention heads, by default 8.
+        qkv_bias : bool, optional
+            Whether to add a bias to the query, key, and value projections, by default False.
+        proj_bias : bool, optional
+            Whether to add a bias to the output projection, by default True.
+        attn_drop : float, optional
+            Dropout rate for the attention weights, by default 0.0.
+        proj_drop : float, optional
+            Dropout rate for the output projection, by default 0.0.
+
+        Raises
+        ------
+        ValueError
+            If `dim` is not divisible by `num_heads`.
+        """
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim {dim} should be divisible by num_heads {num_heads}.")
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of :class:`Attention`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, N, C) where B is the batch size, N is the sequence length, and C is
+            the feature dimension.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (B, N, C) after applying multi-head self-attention.
+        """
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        attn = q @ k.transpose(-2, -1)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    """Memory-efficient multi-head self-attention module using xFormers.
+
+    Parameters
+    ----------
+    dim : int
+        Dimension of the input features.
+    num_heads : int, optional
+        Number of attention heads, by default 8.
+    qkv_bias : bool, optional
+        Whether to add a bias to the query, key, and value projections, by default False.
+    proj_bias : bool, optional
+        Whether to add a bias to the output projection, by default True.
+    attn_drop : float, optional
+        Dropout rate for the attention weights, by default 0.0.
+    proj_drop : float, optional
+        Dropout rate for the output projection, by default 0.0.
+
+    Raises
+    ------
+    ValueError
+        If `dim` is not divisible by `num_heads`.
+    """
+
+    def forward(self, x: torch.Tensor, attn_bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Forward pass of :class:`MemEffAttention`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, N, C) where B is the batch size, N is the sequence length, and C is
+            the feature dimension.
+        attn_bias : Optional[torch.Tensor], optional
+            Attention bias tensor for memory-efficient attention, by default None.
+
+        Raises
+        ------
+        AssertionError
+            If xFormers is not available and `attn_bias` is provided.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (B, N, C) after applying memory-efficient multi-head self-attention.
+        """
+        if not XFORMERS_AVAILABLE:
+            if attn_bias is not None:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

config.json

@@ -0,0 +1,17 @@
+{
+    "model_type": "tapct",
+    "model_size": "small",
+    "model_variant": "2d",
+    "img_size": [224, 224],
+    "patch_size": [16, 16],
+    "in_chans": 1,
+    "num_register_tokens": 4,
+    "init_values": 1e-5,
+    "block_chunks": 0,
+    "architectures": ["TAPCTModel"],
+    "auto_map": {
+        "AutoConfig": "configuration_tapct.TAPCTConfig",
+        "AutoModel": "modeling_tapct.TAPCTModel"
+    }
+}
+

configuration_tapct.py

@@ -0,0 +1,62 @@
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Literal
+from transformers import PretrainedConfig
+
+class TAPCTConfig(PretrainedConfig):
+    """
+    Configuration class for TAP-CT models.
+    
+    Parameters
+    ----------
+    model_size : Literal['small', 'base'], default='base'
+        Size of the model ('small' or 'base')
+    model_variant : Literal['2d', '2.5d', '3d'], default='3d'
+        Variant of the model ('2d', '2.5d', or '3d')
+    img_size : int | tuple | list, default=224
+        Input image size. For 2D: int or tuple[int, int], for 3D: tuple[int, int, int]
+    patch_size : int | tuple | list, default=16
+        Patch size. For 2D: int or tuple[int, int], for 3D: tuple[int, int, int]
+    in_chans : int, default=1
+        Number of input channels (default: 1 for CT scans)
+    num_register_tokens : int, default=4
+        Number of register tokens
+    init_values : float | None, default=None
+        Layer scale init values
+    block_chunks : int, default=0
+        Number of block chunks for FSDP
+    """
+    model_type = "tapct"
+
+    def __init__(
+        self,
+        model_size: Literal['small', 'base'] = 'base',
+        model_variant: Literal['2d', '2.5d', '3d'] = '3d',
+        img_size: int | tuple | list = 224,
+        patch_size: int | tuple | list = 16,
+        in_chans: int = 1,
+        num_register_tokens: int = 4,
+        init_values: float | None = None,
+        block_chunks: int = 0,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.model_size = model_size
+        self.model_variant = model_variant
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+        self.num_register_tokens = num_register_tokens
+        self.init_values = init_values
+        self.block_chunks = block_chunks

drop_path.py

@@ -0,0 +1,86 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+import torch
+from torch import nn
+
+
+def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor of shape (B, *) where B is the batch size and * is any number of additional dimensions.
+    drop_prob : float, optional
+        Probability of dropping a path, by default 0.0
+    training : bool, optional
+        Whether the model is in training mode, by default False. If False, no paths are dropped.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor with the same shape as input x, with paths dropped according to drop_prob.
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    Parameters
+    ----------
+    drop_prob : float, optional
+        Probability of dropping a path, by default None. If None, no paths are dropped.
+        If set to 0.0, it behaves like an identity function.
+    """
+
+    def __init__(self, drop_prob: float = 0.0) -> None:
+        """Inits :class:`DropPath`.
+
+        Parameters
+        ----------
+        drop_prob : float, optional
+            Probability of dropping a path, by default 0.0. If None, no paths are dropped.
+            If set to 0.0, it behaves like an identity function.
+        """
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of :class:`DropPath`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, *) where B is the batch size and * is any number of additional dimensions.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor with the same shape as input x, with paths dropped according to drop_prob.
+        """
+        return drop_path(x, self.drop_prob, self.training)

helpers.py

@@ -0,0 +1,70 @@
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+def make_tuple(x: int | tuple, n: int) -> tuple:
+    """Convert an integer or a tuple to an n-tuple.
+
+    Parameters
+    ----------
+    x : int or tuple
+        Input value which can be an integer or a tuple of n integers.
+    n : int
+        The length of the tuple to return.
+
+    Returns
+    -------
+    tuple
+        A tuple of n integers.
+    """
+    if isinstance(x, tuple) or isinstance(x, list):
+        if len(x) != n:
+            raise ValueError(f"Expected a tuple of length {n}, got {len(x)}")
+        if not all(isinstance(i, int) for i in x):
+            raise ValueError("All elements in the tuple must be integers")
+        return tuple(x)
+
+    if not isinstance(x, int):
+        raise TypeError(f"Expected int, got {type(x)}")
+    return (x,) * n
+
+
+def make_2tuple(x: int | tuple[int, int]) -> tuple[int, int]:
+    """Convert an integer or a tuple to a 2-tuple.
+
+    Parameters
+    ----------
+    x : int or tuple[int, int]
+        Input value which can be an integer or a tuple of two integers with two elements.
+
+    Returns
+    -------
+    tuple[int, int]
+        A tuple of two integers.
+    """
+    return make_tuple(x, 2)
+
+
+def make_3tuple(x: int | tuple[int, int, int]) -> tuple[int, int, int]:
+    """Convert an integer or a tuple to a 3-tuple.
+
+    Parameters
+    ----------
+    x : int or tuple[int, int, int]
+        Input value which can be an integer or a tuple of three integers with three elements.
+
+    Returns
+    -------
+    tuple[int, int, int]
+        A tuple of three integers.
+    """
+    return make_tuple(x, 3)

layer_scale.py

@@ -0,0 +1,74 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+from typing import Union
+
+import torch
+from torch import nn
+
+
+class LayerScale(nn.Module):
+    """Layer scale module for scaling the output of a layer.
+
+    Parameters
+    ----------
+    dim : int
+        Dimension of the layer scale.
+    init_values : float or torch.Tensor, optional
+        Initial values for the layer scale, by default 1e-5. If a tensor is provided, it should have shape (dim,).
+    inplace : bool, optional
+        Whether to perform the operation in-place, by default False.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, torch.Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        """Inits :class:`LayerScale
+
+        Parameters
+        ----------
+        dim : int
+            Dimension of the layer scale.
+        init_values : float or torch.Tensor, optional
+            Initial values for the layer scale, by default 1e-5. If a tensor is provided, it should have shape (dim,).
+        inplace : bool, optional
+            Whether to perform the operation in-place, by default False.
+        """
+        super().__init__()
+
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of :class:`LayerScale`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, N, C) where B is the batch size, N is the sequence length, and C is
+            the feature dimension.
+
+        Returns
+        -------
+        torch.Tensor
+            Scaled output tensor of shape (B, N, C).
+        """
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma

mlp.py

@@ -0,0 +1,102 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+from typing import Callable, Optional
+
+import torch
+from torch import nn
+
+
+class Mlp(nn.Module):
+    """Multi-layer perceptron (MLP) module.
+
+    Creates a simple MLP with two linear layers and an activation function in between and dropout after each layer.
+
+    Parameters
+    ----------
+    in_features : int
+        Number of input features.
+    hidden_features : int, optional
+        Number of hidden features, by default 4 * in_features.
+    out_features : int, optional
+        Number of output features, by default in_features.
+    act_layer : Callable[..., nn.Module], optional
+        Activation layer, by default nn.GELU.
+    drop : float, optional
+        Dropout rate, by default 0.0.
+    bias : bool, optional
+        Whether to use bias in the linear layers, by default True.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        """Inits :class:`Mlp`.
+
+        Parameters
+        ----------
+
+        in_features : int
+            Number of input features.
+        hidden_features : int, optional
+            Number of hidden features, by default 4 * in_features.
+        out_features : int, optional
+            Number of output features, by default in_features.
+        act_layer : Callable[..., nn.Module], optional
+            Activation layer, by default nn.GELU.
+        drop : float, optional
+            Dropout rate, by default 0.0.
+        bias : bool, optional
+            Whether to use bias in the linear layers, by default True.
+        """
+        super().__init__()
+
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of :class:`Mlp`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, N, C) where B is the batch size, N is the sequence length, and C is
+            the feature dimension.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (B, N, out_features) after applying the MLP.
+        """
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d6b145206ed34228418d303833412f20da2e98b586c97f0b2139a48025c27851
+size 85937736

modeling_tapct.py

@@ -0,0 +1,155 @@
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Optional
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+import torch
+
+from .configuration_tapct import TAPCTConfig
+from .vision_transformer import vit_small, vit_base
+from .vision_transformer_3d import vit_3d_small, vit_3d_base
+from .vision_transformer_base import DinoVisionTransformerBase
+
+class TAPCTPreTrainedModel(PreTrainedModel):
+    config_class = TAPCTConfig
+    base_model_prefix = "tapct"
+
+
+class TAPCTModel(TAPCTPreTrainedModel):
+    """
+    TAP-CT Vision Transformer model based on DINOv2: https://github.com/facebookresearch/dinov2.
+    
+    This model outputs raw hidden states and does not include any task-specific head.
+    """
+    
+    def __init__(self, config: TAPCTConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        self.model: DinoVisionTransformerBase
+
+        match config.model_variant:
+            case "2d":
+                if config.model_size == "small":
+                    self.model = vit_small(
+                        img_size=config.img_size,
+                        patch_size=config.patch_size,
+                        num_register_tokens=config.num_register_tokens,
+                        in_chans=config.in_chans,
+                        init_values=config.init_values,
+                        block_chunks=config.block_chunks
+                    )
+                elif config.model_size == "base":
+                    self.model = vit_base(
+                        img_size=config.img_size,
+                        patch_size=config.patch_size,
+                        num_register_tokens=config.num_register_tokens,
+                        in_chans=config.in_chans,
+                        init_values=config.init_values,
+                        block_chunks=config.block_chunks
+                    )
+                else:
+                    raise ValueError(f"Model size '{config.model_size}' not supported for 2D")
+                    
+            case "2.5d" | "3d":
+                if config.model_size == "small":
+                    self.model = vit_3d_small(
+                        img_size=config.img_size,
+                        patch_size=config.patch_size,
+                        num_register_tokens=config.num_register_tokens,
+                        in_chans=config.in_chans,
+                        init_values=config.init_values,
+                        block_chunks=config.block_chunks
+                    )
+                elif config.model_size == "base":
+                    self.model = vit_3d_base(
+                        img_size=config.img_size,
+                        patch_size=config.patch_size,
+                        num_register_tokens=config.num_register_tokens,
+                        in_chans=config.in_chans,
+                        init_values=config.init_values,
+                        block_chunks=config.block_chunks
+                    )
+                else:
+                    raise ValueError(f"Model size '{config.model_size}' not supported for 3D")
+                    
+            case _:
+                raise ValueError(f"Model variant '{config.model_variant}' not supported. Use '2d', '2.5d', or '3d'.")
+        
+        # Initialize weights
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> BaseModelOutputWithPooling:
+        """
+        Forward pass of the TAP-CT model.
+        
+        Parameters
+        ----------
+        pixel_values : torch.Tensor
+            Input images. Shape (B, C, H, W) for 2D or (B, C, D, H, W) for 3D
+        output_hidden_states : Optional[bool], optional
+            Whether to return hidden states from all layers
+        return_dict : Optional[bool], optional
+            Whether to return a ModelOutput instead of a plain tuple
+        
+        Returns
+        -------
+        BaseModelOutputWithPooling
+            Contains:
+            - last_hidden_state: Patch token features from the last layer
+            - pooler_output: CLS token from the last layer  
+            - hidden_states: (optional) All hidden states if output_hidden_states=True
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+
+        if output_hidden_states:
+            outputs_tuple = self.model.get_intermediate_layers(
+                pixel_values,
+                n=self.model.n_blocks,
+                return_class_token=True,
+                reshape=False
+            )
+            outputs = tuple(o[0] for o in outputs_tuple)
+            class_tokens = tuple(o[1] for o in outputs_tuple)
+            
+            last_hidden_state = outputs[-1]
+            pooler_output = class_tokens[-1]
+            hidden_states = outputs
+        else:
+            outputs_tuple = self.model.get_intermediate_layers(
+                pixel_values,
+                n=1,
+                return_class_token=True,
+                reshape=False
+            )
+            last_hidden_state = outputs_tuple[0][0]
+            pooler_output = outputs_tuple[0][1]
+            hidden_states = None
+
+        if not return_dict:
+            return tuple(
+                v for v in [last_hidden_state, pooler_output, hidden_states] 
+                if v is not None
+            )
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooler_output,
+            hidden_states=hidden_states
+        )

patch_embed.py

@@ -0,0 +1,272 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+from typing import Callable, Optional
+
+import torch
+from torch import nn
+
+from .helpers import make_2tuple, make_3tuple
+
+
+class PatchEmbed(nn.Module):
+    """Patch embedding layer for Vision Transformers for 2D images.
+
+    This layer divides the input image into patches and projects them into a higher-dimensional space.
+
+    Parameters
+    ----------
+    img_size : int or tuple[int, int], optional
+        Size of the input image. If an integer is provided, it is assumed to be square (img_size, img_size).
+        If a tuple is provided, it should be of the form (height, width), by default 224.
+    patch_size : int or tuple[int, int], optional
+        Size of the patches to be extracted from the input image. If an integer is provided, it is assumed to be square
+        (patch_size, patch_size). If a tuple is provided, it should be of the form (height, width), by default 16.
+    in_chans : int, optional
+        Number of input channels in the image, by default 3 (for RGB images).
+    embed_dim : int, optional
+        Dimension of the embedding space to which the patches will be projected, by default 768.
+    norm_layer : Callable, optional
+        Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
+    flatten_embedding : bool, optional
+        Whether to flatten the embedding output, by default True.
+    """
+
+    def __init__(
+        self,
+        img_size: int | tuple[int, int] = 224,
+        patch_size: int | tuple[int, int] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        """Inits :class:`PatchEmbed`.
+
+        Parameters
+        ----------
+        img_size : int or tuple[int, int], optional
+            Size of the input image. If an integer is provided, it is assumed to be square (img_size, img_size).
+            If a tuple is provided, it should be of the form (height, width), by default 224.
+        patch_size : int or tuple[int, int], optional
+            Size of the patches to be extracted from the input image. If an integer is provided, it is assumed to be square
+            (patch_size, patch_size). If a tuple is provided, it should be of the form (height, width), by default 16.
+        in_chans : int, optional
+            Number of input channels in the image, by default 3 (for RGB images).
+        embed_dim : int, optional
+            Dimension of the embedding space to which the patches will be projected, by default 768.
+        norm_layer : Callable, optional
+            Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
+        flatten_embedding : bool, optional
+            Whether to flatten the embedding output, by default True.
+        """
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of :class:`PatchEmbed`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, C, H, W) where B is the batch size, C is the number of channels,
+            H is the height, and W is the width of the input image.
+
+        Raises
+        ------
+        ValueError
+            If the input image dimensions are not compatible with the patch size.
+        """
+        _, _, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+        if H % patch_H != 0:
+            raise ValueError(f"Input image height {H} is not a multiple of patch height {patch_H}")
+        if W % patch_W != 0:
+            raise ValueError(f"Input image width {W} is not a multiple of patch width: {patch_W}")
+
+        x = self.proj(x)  # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)  # B HW C
+
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
+        return x
+
+    def flops(self) -> float:
+        """Calculate the number of floating point operations (FLOPs) for the patch embedding layer.
+
+        Returns
+        -------
+        float
+            The number of FLOPs for the patch embedding layer.
+        """
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if not isinstance(self.norm, nn.Identity):
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class PatchEmbed3d(nn.Module):
+    """Patch embedding layer for Vision Transformers for 3D images.
+
+    This layer divides the input 3D image volume into patches and projects them into a higher-dimensional space.
+
+    Parameters
+    ----------
+    img_size : int or tuple[int, int, int], optional
+        Size of the input image volume. If an integer is provided, it is assumed to be cubic (img_size, img_size, img_size).
+        If a tuple is provided, it should be of the form (depth, height, width), by default 224.
+    patch_size : int or tuple[int, int, int], optional
+        Size of the patches to be extracted from the input image volume. If an integer is provided, it is assumed to be cubic
+        (patch_size, patch_size, patch_size). If a tuple is provided, it should be of the form (depth, height, width), by default 16.
+    in_chans : int, optional
+        Number of input channels in the image volume, by default 3 (for RGB images).
+    embed_dim : int, optional
+        Dimension of the embedding space to which the patches will be projected, by default 768.
+    norm_layer : Callable, optional
+        Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
+    flatten_embedding : bool, optional
+        Whether to flatten the embedding output, by default True.
+    """
+
+    def __init__(
+        self,
+        img_size: int | tuple[int, int, int] = 224,
+        patch_size: int | tuple[int, int, int] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        """Inits :class:`PatchEmbed3d`.
+
+        Parameters
+        ----------
+        img_size : int or tuple[int, int, int], optional
+            Size of the input image volume. If an integer is provided, it is assumed to be cubic
+            (img_size, img_size, img_size).
+            If a tuple is provided, it should be of the form (depth, height, width), by default 224.
+        patch_size : int or tuple[int, int, int], optional
+            Size of the patches to be extracted from the input image volume. If an integer is provided, it is
+            assumed to be cubic (patch_size, patch_size, patch_size). If a tuple is provided, it should be of the
+            form (depth, height, width), by default 16.
+        in_chans : int, optional
+            Number of input channels in the image volume, by default 3 (for RGB images).
+        embed_dim : int, optional
+            Dimension of the embedding space to which the patches will be projected, by default 768.
+        norm_layer : Callable, optional
+            Normalization layer to apply to the embeddings, by default None. If None, no normalization is applied.
+        flatten_embedding : bool, optional
+            Whether to flatten the embedding output, by default True.
+        """
+        super().__init__()
+
+        image_DHW = make_3tuple(img_size)
+        patch_DHW = make_3tuple(patch_size)
+
+        patch_grid_size = (
+            image_DHW[0] // patch_DHW[0],
+            image_DHW[1] // patch_DHW[1],
+            image_DHW[2] // patch_DHW[2],
+        )
+
+        self.img_size = image_DHW
+        self.patch_size = patch_DHW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1] * patch_grid_size[2]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_DHW, stride=patch_DHW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of :class:`PatchEmbed3d`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, C, D, H, W) where B is the batch size, C is the number of channels,
+            D is the depth, H is the height, and W is the width of the input volume.
+
+        Raises
+        ------
+        ValueError
+            If the input volume dimensions are not compatible with the patch size.
+        """
+        _, _, D, H, W = x.shape
+        patch_D, patch_H, patch_W = self.patch_size
+        if D % patch_D != 0:
+            raise ValueError(f"Input volume depth {D} is not a multiple of patch depth {patch_D}")
+        if H % patch_H != 0:
+            raise ValueError(f"Input volume height {H} is not a multiple of patch height {patch_H}")
+        if W % patch_W != 0:
+            raise ValueError(f"Input volume width {W} is not a multiple of patch width {patch_W}")
+
+        x = self.proj(x)  # B C D H W
+        D, H, W = x.size(2), x.size(3), x.size(4)
+        x = x.flatten(2).transpose(1, 2)  # B (DHW) C
+
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, D, H, W, self.embed_dim)  # B D H W C
+        return x
+
+    def flops(self) -> float:
+        """Calculate the number of floating point operations (FLOPs) for the patch embedding 3D layer.
+
+        Returns
+        -------
+        float
+            The number of FLOPs for the patch embedding layer.
+        """
+        Do, Ho, Wo = self.patches_resolution
+        flops = (
+            Do
+            * Ho
+            * Wo
+            * self.embed_dim
+            * self.in_chans
+            * (self.patch_size[0] * self.patch_size[1] * self.patch_size[2])
+        )
+        if not isinstance(self.norm, nn.Identity):
+            flops += Do * Ho * Wo * self.embed_dim
+        return flops

swiglu_ffn.py

@@ -0,0 +1,191 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+import os
+import warnings
+from typing import Callable, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class SwiGLUFFN(nn.Module):
+    r"""SwiGLU Feed-Forward Network (FFN) layer.
+
+    SwiGLU Feed-Forward Network (FFN) layer.
+
+    This module applies a two-layer position-wise feed-forward transformation with a SwiGLU activation: 
+    a gated unit combining the SiLU nonlinearity with an elementwise multiplication.
+
+    Given input tensor ``x`` of shape ``(B, d)``, the computation is:
+
+    .. math::
+
+        [z_1, z_2] = x W_{12} + b_{12} \\\\
+        h = \mathrm{SiLU}(z_1) \odot z_2 \\\\
+        y = h W_3 + b_3
+
+    where:
+        - :math:`W_{12} \in \mathbb{R}^{d \times 2h}`, :math:`b_{12} \in \mathbb{R}^{2h}`
+        - :math:`W_3 \in \mathbb{R}^{h \times d_{\text{out}}}`, :math:`b_3 \in \mathbb{R}^{d_{\text{out}}}`
+        - :math:`\mathrm{SiLU}(x) = x \cdot \sigma(x)` is the Sigmoid Linear Unit
+        - :math:`\odot` denotes elementwise multiplication
+
+    Parameters
+    ----------
+    in_features : int
+        Input feature dimensionality (d).
+    hidden_features : int, optional
+        Hidden layer dimensionality (h). Defaults to in_features.
+    out_features : int, optional
+        Output feature dimensionality (d_out). Defaults to in_features.
+    act_layer : Callable[..., nn.Module], optional
+        Unused. Included for compatibility.
+    drop : float, optional
+        Dropout rate (unused).
+    bias : bool, optional
+        Whether to include bias terms in linear layers.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        """Inits :class:`SwiGLUFFN`.
+
+        Parameters
+        ----------
+        in_features : int
+            Input feature dimensionality (d).
+        hidden_features : int, optional
+            Hidden layer dimensionality (h). Defaults to in_features.
+        out_features : int, optional
+            Output feature dimensionality (d_out). Defaults to in_features.
+        act_layer : Callable[..., nn.Module], optional
+            Unused. Included for compatibility.
+        drop : float, optional
+            Dropout rate (unused).
+        bias : bool, optional
+            Whether to include bias terms in linear layers.
+        """
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of :class:`SwiGLUFFN`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, N, C) where B is the batch size, N is the sequence length, and C is
+            the input feature dimension.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (B, N, out_features) after applying the SwiGLU feed-forward network.
+        """
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import SwiGLU
+
+        XFORMERS_AVAILABLE = True
+        warnings.warn("xFormers is available (SwiGLU)")
+    else:
+        warnings.warn("xFormers is disabled (SwiGLU)")
+        raise ImportError
+except ImportError:
+    SwiGLU = SwiGLUFFN
+    XFORMERS_AVAILABLE = False
+
+    warnings.warn("xFormers is not available (SwiGLU)")
+
+
+class SwiGLUFFNFused(SwiGLU):
+    """Fused SwiGLU Feed-Forward Network (FFN) layer.
+
+    Fused SwiGLU Feed-Forward Network (FFN) layer that uses xFormers' fused implementation if available.
+    This layer combines the linear transformations and activation into a single operation for improved performance.
+
+    Parameters
+    ----------
+    in_features : int
+        Input feature dimensionality (d).
+    hidden_features : int, optional
+        Hidden layer dimensionality (h). Defaults to in_features.
+    out_features : int, optional
+        Output feature dimensionality (d_out). Defaults to in_features.
+    act_layer : Callable[..., nn.Module], optional
+        Unused. Included for compatibility.
+    drop : float, optional
+        Dropout rate (unused).
+    bias : bool, optional
+        Whether to include bias terms in linear layers.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        """Inits :class:`SwiGLUFFNF
+
+        Parameters
+        ----------
+        in_features : int
+            Input feature dimensionality (d).
+        hidden_features : int, optional
+            Hidden layer dimensionality (h). Defaults to in_features.
+        out_features : int, optional
+            Output feature dimensionality (d_out). Defaults to in_features.
+        act_layer : Callable[..., nn.Module], optional
+            Unused. Included for compatibility.
+        drop : float, optional
+            Dropout rate (unused).
+        bias : bool, optional
+            Whether to include bias terms in linear layers.
+        """
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+        super().__init__(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            out_features=out_features,
+            bias=bias,
+        )

transformer_block.py

@@ -0,0 +1,565 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+import os
+import warnings
+from typing import Any, Callable, Dict, Optional, Tuple
+
+import torch
+from torch import nn
+
+from .attention import Attention, MemEffAttention
+from .drop_path import DropPath
+from .layer_scale import LayerScale
+from .mlp import Mlp
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import fmha, index_select_cat, scaled_index_add
+
+        XFORMERS_AVAILABLE = True
+        warnings.warn("xFormers is available (Block)")
+    else:
+        warnings.warn("xFormers is disabled (Block)")
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+
+    warnings.warn("xFormers is not available (Block)")
+
+
+class Block(nn.Module):
+    """Transformer block with multi-head self-attention and MLP.
+
+    Parameters
+    ----------
+    dim : int
+        Dimension of the input features.
+    num_heads : int
+        Number of attention heads, by default 8.
+    mlp_ratio : float, optional
+        Ratio of the hidden dimension in the MLP to the input dimension, by default 4.0.
+    qkv_bias : bool, optional
+        Whether to add a bias to the query, key, and value projections, by default False.
+    proj_bias : bool, optional
+        Whether to add a bias to the output projection, by default True.
+    ffn_bias : bool, optional
+        Whether to add a bias to the MLP layers, by default True.
+    drop : float, optional
+        Dropout rate for the MLP layers, by default 0.0.
+    attn_drop : float, optional
+        Dropout rate for the attention weights, by default 0.0.
+    init_values : float or torch.Tensor, optional
+        Initial values for the layer scale, by default None. If a tensor is provided, it should have shape (dim,).
+    drop_path : float, optional
+        Drop path rate for stochastic depth, by default 0.0.
+    act_layer : Callable[..., nn.Module], optional
+        Activation layer for the MLP, by default nn.GELU.
+    norm_layer : Callable[..., nn.Module], optional
+        Normalization layer, by default nn.LayerNorm.
+    attn_class : Callable[..., nn.Module], optional
+        Attention class to use, by default Attention. Can be replaced with :class:`MemEffAttention` for memory-efficient
+        attention.
+    ffn_layer : Callable[..., nn.Module], optional
+        MLP class to use, by default Mlp.
+
+    Raises
+    ------
+    ValueError
+        If `dim` is not divisible by `num_heads`.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = Attention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+    ) -> None:
+        """Inits :class:`Block`.
+
+        Parameters
+        ----------
+        dim : int
+            Dimension of the input features.
+        num_heads : int
+            Number of attention heads, by default 8.
+        mlp_ratio : float, optional
+            Ratio of the hidden dimension in the MLP to the input dimension, by default 4.0.
+        qkv_bias : bool, optional
+            Whether to add a bias to the query, key, and value projections, by default False.
+        proj_bias : bool, optional
+            Whether to add a bias to the output projection, by default True.
+        ffn_bias : bool, optional
+            Whether to add a bias to the MLP layers, by default True.
+        drop : float, optional
+            Dropout rate for the MLP layers, by default 0.0.
+        attn_drop : float, optional
+            Dropout rate for the attention weights, by default 0.0.
+        init_values : float or torch.Tensor, optional
+            Initial values for the layer scale, by default None. If a tensor is provided, it should have shape (dim,).
+        drop_path : float, optional
+            Drop path rate for stochastic depth, by default 0.0.
+        act_layer : Callable[..., nn.Module], optional
+            Activation layer for the MLP, by default nn.GELU.
+        norm_layer : Callable[..., nn.Module], optional
+            Normalization layer, by default nn.LayerNorm.
+        attn_class : Callable[..., nn.Module], optional
+            Attention class to use, by default Attention. Can be replaced with :class:`MemEffAttention` for
+            memory-efficient attention.
+        ffn_layer : Callable[..., nn.Module], optional
+            MLP class to use, by default Mlp.
+
+        Raises
+        ------
+        ValueError
+            If `dim` is not divisible by `num_heads`.
+        """
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim {dim} should be divisible by num_heads {num_heads}.")
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of :class:`Block`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, N, C) where B is the batch size, N is the sequence length, and C is
+            the feature dimension.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (B, N, C) after applying the transformer block.
+        """
+
+        def attn_residual_func(x: torch.Tensor) -> torch.Tensor:
+            return self.ls1(self.attn(self.norm1(x)))
+
+        def ffn_residual_func(x: torch.Tensor) -> torch.Tensor:
+            return self.ls2(self.mlp(self.norm2(x)))
+
+        if self.training and self.sample_drop_ratio > 0.1:
+            # the overhead is compensated only for a drop path rate larger than 0.1
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+        elif self.training and self.sample_drop_ratio > 0.0:
+            x = x + self.drop_path1(attn_residual_func(x))
+            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
+        else:
+            x = x + attn_residual_func(x)
+            x = x + ffn_residual_func(x)
+        return x
+
+
+def drop_add_residual_stochastic_depth(
+    x: torch.Tensor,
+    residual_func: Callable[[torch.Tensor], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> torch.Tensor:
+    """Applies stochastic depth by dropping a subset of samples in the batch and adding a residual.
+
+    This function extracts a random subset of the batch, applies a residual function to it, and adds the result back
+    to the original tensor, scaling the residual appropriately.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor of shape (B, N, D) where B is the batch size, N is the sequence length, and D is the
+        feature dimension.
+    residual_func : Callable[[torch.Tensor], torch.Tensor]
+        Function that takes a tensor of shape (B', N, D) and returns a tensor of the same shape, representing the
+        residual.
+    sample_drop_ratio : float, optional
+        Ratio of samples to drop from the batch, by default 0.0. If set to 0.0, no samples are dropped.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor of the same shape as input x, with the residual added back to the original tensor.
+    """
+    # 1) extract subset using permutation
+    B = x.shape[0]
+    sample_subset_size = max(int(B * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(B, device=x.device))[:sample_subset_size]
+    x_subset = x[brange]
+
+    # 2) apply residual_func to get residual
+    residual = residual_func(x_subset)
+
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+
+    residual_scale_factor = B / sample_subset_size
+
+    # 3) add the residual
+    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x: torch.Tensor, sample_drop_ratio: float = 0.0) -> tuple[torch.Tensor, float]:
+    """Generates random indices for dropping samples in the batch and computes the scale factor for the residual.
+
+    This function extracts a random subset of the batch and computes a scale factor based on the original batch size
+    and the size of the subset. The scale factor is used to scale the residual when it is added back to the original
+    tensor.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor of shape (B, N, D) where B is the batch size, N is the sequence length, and D is the
+        feature dimension.
+    sample_drop_ratio : float, optional
+        Ratio of samples to drop from the batch, by default 0.0. If set to 0.0, no samples are dropped.
+
+    Returns
+    -------
+    tuple[torch.Tensor, float]
+        A tuple containing:
+        - brange: A tensor of indices representing the subset of the batch to keep.
+        - residual_scale_factor: A float representing the scale factor for the residual.
+    """
+
+    B = x.shape[0]
+    sample_subset_size = max(int(B * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(B, device=x.device))[:sample_subset_size]
+    residual_scale_factor = B / sample_subset_size
+    return brange, residual_scale_factor
+
+
+def add_residual(
+    x: torch.Tensor,
+    brange: torch.Tensor,
+    residual: torch.Tensor,
+    residual_scale_factor: float,
+    scaling_vector: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Adds a residual to the input tensor, scaling it appropriately.
+
+    This function takes a tensor `x`, a set of indices `brange`, and a residual tensor, and adds the residual to the
+    corresponding indices in `x`. If a scaling vector is provided, it scales the residual before adding it.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor of shape (B, N, D) where B is the batch size, N is the sequence length, and D is the
+        feature dimension.
+    brange : torch.Tensor
+        torch.Tensor of indices representing the subset of the batch to which the residual will be added.
+    residual : torch.Tensor
+        Residual tensor of shape (B', N, D) where B' is the size of the subset defined by `brange`.
+    residual_scale_factor : float
+        Scale factor for the residual, computed as the ratio of the original batch size to the subset size.
+    scaling_vector : Optional[torch.Tensor], optional
+        Scaling vector to scale the residual before adding it, by default None. If provided, it should have shape (D,).
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor of the same shape as input `x`, with the residual added back to the original tensor.
+    """
+    if scaling_vector is None:
+        x_flat = x.flatten(1)
+        residual = residual.flatten(1)
+        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    else:
+        x_plus_residual = scaled_index_add(
+            x,
+            brange,
+            residual.to(dtype=x.dtype),
+            scaling=scaling_vector,
+            alpha=residual_scale_factor,
+        )
+    return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(
+    x_list: list[torch.Tensor], branges: Optional[list[torch.Tensor]] = None
+) -> tuple[Any, torch.Tensor]:
+    """Get attention bias and concatenate tensors from a list of tensors.
+
+    This function checks if the attention bias for the given shapes is already cached. If not, it creates a new
+    attention bias using the `fmha.BlockDiagonalMask` from xFormers. It then concatenates the tensors in `x_list`
+    based on the provided `branges`. If `branges` is not provided, it concatenates the tensors directly.
+
+    Parameters
+    ----------
+    x_list : list of torch.Tensors
+        List of tensors to concatenate. Each tensor should have shape (B, N, D) where B is the batch size, N is the
+        sequence length, and D is the feature dimension.
+    branges : list of torch.Tensors, optional
+        List of tensors containing indices for selecting samples from the batch. If provided, it will index select
+        and concatenate the tensors in `x_list`. If not provided, it will concatenate the tensors directly.
+
+    Returns
+    -------
+    tuple[Any, torch.Tensor]
+        A tuple containing:
+        - attn_bias: Attention bias tensor created using `fmha.BlockDiagonalMask` from xFormers.
+        - cat_tensors: Concatenated tensor of shape (1, B', D) where B' is the total number of samples selected from
+        the batch based on `branges` or the total number of samples in `x_list` if `branges` is not provided.
+    If `branges` is provided, the concatenated tensor will have shape (1, sum of sizes in branges, D).
+    If `branges` is not provided, the concatenated tensor will have shape (1, sum of batch sizes in x_list, D).
+    """
+
+    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+    if all_shapes not in attn_bias_cache.keys():
+        seqlens = []
+        for b, x in zip(batch_sizes, x_list):
+            for _ in range(b):
+                seqlens.append(x.shape[1])
+        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+        attn_bias._batch_sizes = batch_sizes
+        attn_bias_cache[all_shapes] = attn_bias
+
+    if branges is not None:
+        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+    else:
+        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+        cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+    return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: list[torch.Tensor],
+    residual_func: Callable[[torch.Tensor, Any], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> list[torch.Tensor]:
+    """Applies stochastic depth to a list of tensors, dropping a subset of samples in each tensor and adding a residual.
+    This function processes a list of tensors, generating random indices for dropping samples in each tensor,
+    computing the attention bias, and applying a residual function to each tensor. The results are then combined
+    and returned as a list of tensors.
+
+    Parameters
+    ----------
+    x_list : list of torch.Tensors
+        List of tensors to process. Each tensor should have shape (B, N, D) where B is the batch size, N is the sequence
+        length, and D is the feature dimension.
+    residual_func : Callable[[torch.Tensor, Any], torch.Tensor]
+        Function that takes a tensor of shape (B', N, D) and an attention bias (if applicable) and returns a tensor of
+        the same shape, representing the residual.
+    sample_drop_ratio : float, optional
+        Ratio of samples to drop from the batch, by default 0.0. If set to 0.0, no samples are dropped.
+    scaling_vector : Optional[torch.Tensor], optional
+        Scaling vector to scale the residual before adding it, by default None. If provided, it should have shape (D,).
+
+    Returns
+    -------
+    list of torch.Tensors
+        List of output tensors, each of the same shape as the corresponding input tensor in `x_list`, with the residual
+        added back to the original tensor.
+    """
+    # 1) generate random set of indices for dropping samples in the batch
+    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+    branges = [s[0] for s in branges_scales]
+    residual_scale_factors = [s[1] for s in branges_scales]
+
+    # 2) get attention bias and index+concat the tensors
+    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+    # 3) apply residual_func to get residual, and split the result
+    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+    outputs = []
+    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+    return outputs
+
+
+class NestedTensorBlock(Block):
+    """Transformer block with multi-head self-attention and MLP, supporting nested tensors.
+
+    This class extends the :class:`Block` class to support nested tensors, allowing for more flexible input shapes.
+
+    Parameters
+    ----------
+    dim : int
+        Dimension of the input features.
+    num_heads : int
+        Number of attention heads, by default 8.
+    mlp_ratio : float, optional
+        Ratio of the hidden dimension in the MLP to the input dimension, by default 4.0.
+    qkv_bias : bool, optional
+        Whether to add a bias to the query, key, and value projections, by default False.
+    proj_bias : bool, optional
+        Whether to add a bias to the output projection, by default True.
+    ffn_bias : bool, optional
+        Whether to add a bias to the feed-forward network, by default True.
+    drop : float, optional
+        Dropout rate for the MLP layers, by default 0.0.
+    attn_drop : float, optional
+        Dropout rate for the attention weights, by default 0.0.
+    init_values : float or torch.Tensor, optional
+        Initial values for the layer scale, by default None. If a tensor is provided, it should have shape (dim,).
+    drop_path : float, optional
+        Drop path rate for stochastic depth, by default 0.0.
+    act_layer : Callable[..., nn.Module], optional
+        Activation layer for the MLP, by default nn.GELU.
+    norm_layer : Callable[..., nn.Module], optional
+        Normalization layer, by default nn.LayerNorm.
+    attn_class : Callable[..., nn.Module], optional
+        Attention class to use, by default Attention. Can be replaced with :class:`MemEffAttention` for
+        memory-efficient attention.
+    ffn_layer : Callable[..., nn.Module], optional
+        MLP class to use, by default :class:`Mlp`.
+    sample_drop_ratio : float, optional
+        Drop path rate for stochastic depth, by default 0.0. This is used to control the stochastic depth
+        during training.
+    """
+
+    def forward_nested(self, x_list: list[torch.Tensor]) -> list[torch.Tensor]:
+        """Forward pass for list of tensors, applying attention and MLP with stochastic depth.
+
+        This method applies the attention and MLP layers to a list of tensors, applying stochastic depth if the model is
+        in training mode and `sample_drop_ratio` is greater than 0.0. It uses the :class:`MemEffAttention` class
+        for memory-efficient attention. The method expects `x_list` to be a list of tensors, where each tensor has
+        the same feature dimension. If the model is not in training mode or `sample_drop_ratio` is 0.0,
+        it applies the attention and MLP layers without stochastic depth.
+
+        Parameters
+        ----------
+        x_list : list[torch.Tensor]
+            List of tensors to process. Each tensor should have shape (B, N, D) where B is the batch size, N is the
+            sequence length, and D is the feature dimension.
+
+        Returns
+        -------
+        list[torch.Tensor]
+            List of processed tensors, each with the same shape as the corresponding input tensor in `x_list`.
+        """
+        assert isinstance(self.attn, MemEffAttention)
+
+        if self.training and self.sample_drop_ratio > 0.0:
+
+            def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+            def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.mlp(self.norm2(x))
+
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=(self.ls1.gamma if isinstance(self.ls1, LayerScale) else None),
+            )
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=(self.ls2.gamma if isinstance(self.ls1, LayerScale) else None),
+            )
+            return x_list
+        else:
+
+            def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+            def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.ls2(self.mlp(self.norm2(x)))
+
+            attn_bias, x = get_attn_bias_and_cat(x_list)
+            x = x + attn_residual_func(x, attn_bias=attn_bias)
+            x = x + ffn_residual_func(x)
+            return attn_bias.split(x)
+
+    def forward(self, x_or_x_list: torch.Tensor | list[torch.Tensor]) -> torch.Tensor | list[torch.Tensor]:
+        """Forward pass of :class:`NestedTensorBlock`.
+
+        Parameters
+        ----------
+        x_or_x_list : torch.Tensor or list[torch.Tensor]
+            Input tensor or list of tensors. If a tensor is provided, it should have shape (B, N, D) where B is the
+            batch size, N is the sequence length, and D is the feature dimension. If a list of tensors is provided,
+            each tensor should have the same shape.
+
+        Returns
+        -------
+        torch.Tensor or list[torch.Tensor]
+            Output tensor or list of tensors after applying the transformer block. If a tensor is provided, the output
+            will be a tensor of the same shape. If a list of tensors is provided, the output will be a list of tensors,
+            each with the same shape as the corresponding input tensor.
+
+        Raises
+        ------
+        AssertionError
+            If `xFormers` is not available.
+        ValueError
+            If `x_or_x_list` is neither a torch.Tensor nor a list of torch.Tensors.
+        """
+        if isinstance(x_or_x_list, torch.Tensor):
+            return super().forward(x_or_x_list)
+        elif isinstance(x_or_x_list, list):
+            if not XFORMERS_AVAILABLE:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise ValueError(
+                f"Expected input to be a torch.Tensor or a list of torch.Tensors, got {type(x_or_x_list)}."
+            )

vision_transformer.py

@@ -0,0 +1,351 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+from typing import Callable
+from typing_extensions import override
+
+import torch
+from torch import nn
+
+from .attention import MemEffAttention
+from .transformer_block import NestedTensorBlock as Block
+from .vision_transformer_base import (
+    DinoVisionTransformerBase,
+    DinoVisionTransformerDim,
+    DinoVisionTransformerFFNLayer,
+)
+
+
+class DinoVisionTransformer(DinoVisionTransformerBase):
+    """DinoVisionTransformer for 2D images.
+
+    Parameters
+    ----------
+    img_size : int or tuple[int, int]
+        Input image size, either a single integer or a tuple of two integers (height, width).
+    patch_size : int or tuple[int, int]
+        Patch size, either a single integer or a tuple of two integers (height, width).
+    in_chans : int
+        Number of input channels, default is 3.
+    embed_dim : int
+        Embedding dimension.
+    depth : int
+        Depth of transformer.
+    num_heads : int
+        Number of attention heads.
+    mlp_ratio : int
+        Ratio of mlp hidden dim to embedding dim.
+    qkv_bias : bool
+        Enable bias for qkv if True.
+    proj_bias : bool
+        Enable bias for proj in attn if True.
+    ffn_bias : bool
+        Enable bias for ffn if True.
+    drop_path_rate : float
+        Stochastic depth rate.
+    drop_path_uniform : bool
+        Apply uniform drop rate across blocks.
+    weight_init : str
+        Weight init scheme.
+    init_values : float
+        Layer-scale init values.
+    act_layer : nn.Module
+        MLP activation layer.
+    block_fn : nn.Module
+        Transformer block class.
+    ffn_layer : DinoVisionTransformerFFNLayer
+        Type of FFN layer to use, can be DinoVisionTransformerFFNLayer.MLP,
+        DinoVisionTransformerFFNLayer.SWIGLU, DinoVisionTransformerFFNLayer.SWIGLU_FUSED,
+        or DinoVisionTransformerFFNLayer.IDENTITY. Default is DinoVisionTransformerFFNLayer.MLP.
+    block_chunks : int
+        Split block sequence into block_chunks units for FSDP wrap.
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 0.
+    interpolate_antialias : str
+        Flag to apply anti-aliasing when interpolating positional embeddings.
+    interpolate_offset : float
+        Work-around offset to apply when interpolating positional embeddings.
+    """
+
+    def __init__(
+        self,
+        img_size: int | tuple[int, int] = 224,
+        patch_size: int | tuple[int, int] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        ffn_bias: bool = True,
+        proj_bias: bool = True,
+        drop_path_rate: float = 0.0,
+        drop_path_uniform: bool = False,
+        init_values: float | None = None,  # for layerscale: None or 0 => no layerscale
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        block_fn: Callable[..., Block] = Block,
+        ffn_layer: DinoVisionTransformerFFNLayer = DinoVisionTransformerFFNLayer.MLP,
+        block_chunks: int = 1,
+        num_register_tokens: int = 0,
+        interpolate_antialias: bool = False,
+        interpolate_offset: float = 0.1,
+    ) -> None:
+        """Inits :class:`DinoVisionTransformer`.
+
+        Parameters
+        ----------
+        img_size : int or tuple[int, int]
+            Input image size, either a single integer or a tuple of two integers (height, width).
+        patch_size : int or tuple[int, int]
+            Patch size, either a single integer or a tuple of two integers (height, width).
+        in_chans : int
+            Number of input channels, default is 3.
+        embed_dim : int
+            Embedding dimension.
+        depth : int
+            Depth of transformer.
+        num_heads : int
+            Number of attention heads.
+        mlp_ratio : int
+            Ratio of mlp hidden dim to embedding dim.
+        qkv_bias : bool
+            Enable bias for qkv if True.
+        proj_bias : bool
+            Enable bias for proj in attn if True.
+        ffn_bias : bool
+            Enable bias for ffn if True.
+        drop_path_rate : float
+            Stochastic depth rate.
+        drop_path_uniform : bool
+            Apply uniform drop rate across blocks.
+        weight_init : str
+            Weight init scheme.
+        init_values : float
+            Layer-scale init values.
+        act_layer : nn.Module
+            MLP activation layer.
+        block_fn : nn.Module
+            Transformer block class.
+        ffn_layer : DinoVisionTransformerFFNLayer
+            Type of FFN layer to use, can be DinoVisionTransformerFFNLayer.MLP,
+            DinoVisionTransformerFFNLayer.SWIGLU, DinoVisionTransformerFFNLayer.SWIGLU_FUSED,
+            or DinoVisionTransformerFFNLayer.IDENTITY. Default is DinoVisionTransformerFFNLayer.MLP.
+        block_chunks : int
+            Split block sequence into block_chunks units for FSDP wrap.
+        num_register_tokens : int
+            Number of extra tokens for the model to deposit information (so-called "registers"). Default is 0.
+        interpolate_antialias : str
+            Flag to apply anti-aliasing when interpolating positional embeddings.
+        interpolate_offset : float
+            Work-around offset to apply when interpolating positional embeddings.
+        """
+        super().__init__(
+            dim=DinoVisionTransformerDim.TWO_D,
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            depth=depth,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            ffn_bias=ffn_bias,
+            proj_bias=proj_bias,
+            drop_path_rate=drop_path_rate,
+            drop_path_uniform=drop_path_uniform,
+            init_values=init_values,
+            act_layer=act_layer,
+            block_fn=block_fn,
+            ffn_layer=ffn_layer,
+            block_chunks=block_chunks,
+            num_register_tokens=num_register_tokens,
+            interpolate_antialias=interpolate_antialias,
+            interpolate_offset=interpolate_offset,
+        )
+
+    @override
+    def _interpolate_and_reshape_pos_embed(
+        self, patch_pos_embed: torch.Tensor, patches_resolution: tuple[int, int], dim: int, interpolation_kwargs: dict
+    ) -> torch.Tensor:
+        """Interpolate and reshape 2D patch positional embeddings.
+
+        Parameters
+        ----------
+        patch_pos_embed : torch.Tensor
+            Positional embedding tensor of shape (1, N, C).
+        patches_resolution : tuple of ints
+            Number of patches along each spatial dimension.
+        dim : int
+            Embedding dimension.
+        interpolation_kwargs : dict
+            Arguments passed to `F.interpolate`.
+
+        Returns
+        -------
+        torch.Tensor
+            Reshaped and interpolated tensor of shape (1, H, W, C),
+            where H, W are the number of patches along height and width.
+        """
+        patch_pos_embed = patch_pos_embed.reshape(1, *patches_resolution, dim).permute(0, 3, 1, 2)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed,
+            mode="bicubic",
+            antialias=self.interpolate_antialias,
+            **interpolation_kwargs,
+        )
+        return patch_pos_embed.permute(0, 2, 3, 1)
+
+
+def vit_small(
+    patch_size: int | tuple[int, int] = 16,
+    num_register_tokens: int = 0,
+    **kwargs,
+) -> DinoVisionTransformer:
+    """Builds a small 2d vision transformer with 384-dimensional embeddings, 12 layers, 6 heads, and 4x MLP ratio.
+
+    Parameters
+    ----------
+    patch_size : int or tuple[int, int]
+        Patch size, either a single integer or a tuple of two integers (height, width). Default is 16.
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 0.
+    kwargs : dict
+        Additional keyword arguments to pass to the :class:`DinoVisionTransformer` constructor.
+
+    Returns
+    -------
+    DinoVisionTransformer
+        A small 2d vision transformer.
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_base(
+    patch_size: int | tuple[int, int] = 16,
+    num_register_tokens: int = 0,
+    **kwargs,
+) -> DinoVisionTransformer:
+    """Builds a base 2d vision transformer with 768-dimensional embeddings, 12 layers, 12 heads, and 4x MLP ratio.
+
+    Parameters
+    ----------
+    patch_size : int or tuple[int, int]
+        Patch size, either a single integer or a tuple of two integers (height, width). Default is 16.
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 0.
+    kwargs : dict
+        Additional keyword arguments to pass to the :class:`DinoVisionTransformer` constructor.
+
+    Returns
+    -------
+    DinoVisionTransformer
+        A base 2d vision transformer.
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_large(
+    patch_size: int | tuple[int, int] = 16,
+    num_register_tokens: int = 0,
+    **kwargs,
+) -> DinoVisionTransformer:
+    """Builds a large 2d vision transformer with 1024-dimensional embeddings, 24 layers, 16 heads, and 4x MLP ratio.
+
+    Parameters
+    ----------
+    patch_size : int or tuple[int, int]
+        Patch size, either a single integer or a tuple of two integers (height, width). Default is 16.
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 0.
+    kwargs : dict
+        Additional keyword arguments to pass to the :class:`DinoVisionTransformer` constructor.
+
+    Returns
+    -------
+    DinoVisionTransformer
+        A large 2d vision transformer.
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_giant2(
+    patch_size: int | tuple[int, int] = 16,
+    num_register_tokens: int = 0,
+    **kwargs,
+) -> DinoVisionTransformer:
+    """Builds a giant2 vision transformer with 1536-dimensional embeddings, 40 layers, 24 heads, and 4x MLP ratio.
+
+    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+
+    Parameters
+    ----------
+    patch_size : int or tuple[int, int]
+        Patch size, either a single integer or a tuple of two integers (height, width). Default is 16.
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 0.
+    kwargs : dict
+        Additional keyword arguments to pass to the :class:`DinoVisionTransformer` constructor.
+
+    Returns
+    -------
+    DinoVisionTransformer
+        A giant2 vision transformer.
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1536,
+        depth=40,
+        num_heads=24,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model

vision_transformer_3d.py

@@ -0,0 +1,322 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+from typing import Callable
+from typing_extensions import override
+
+import torch
+from torch import nn
+
+from .attention import MemEffAttention
+from .transformer_block import NestedTensorBlock as Block
+from .vision_transformer_base import (
+    DinoVisionTransformerBase,
+    DinoVisionTransformerDim,
+    DinoVisionTransformerFFNLayer,
+)
+
+
+class DinoVisionTransformer3d(DinoVisionTransformerBase):
+    """DinoVisionTransformer for 3D images.
+
+    Parameters
+    ----------
+    img_size : int or tuple[int, int, int]
+        Input image size, either a single integer or a tuple of three integers (depth, height, width).
+    patch_size : int or tuple[int, int, int]
+        Patch size, either a single integer or a tuple of three integers (depth, height, width).
+    in_chans : int
+        Number of input channels, default is 3.
+    embed_dim : int
+        Embedding dimension.
+    depth : int
+        Depth of transformer.
+    num_heads : int
+        Number of attention heads.
+    mlp_ratio : int
+        Ratio of mlp hidden dim to embedding dim.
+    qkv_bias : bool
+        Enable bias for qkv if True.
+    proj_bias : bool
+        Enable bias for proj in attn if True.
+    ffn_bias : bool
+        Enable bias for ffn if True.
+    drop_path_rate : float
+        Stochastic depth rate.
+    drop_path_uniform : bool
+        Apply uniform drop rate across blocks.
+    weight_init : str
+        Weight init scheme.
+    init_values : float
+        Layer-scale init values.
+    act_layer : nn.Module
+        MLP activation layer.
+    block_fn : nn.Module
+        Transformer block class.
+    ffn_layer : DinoVisionTransformerFFNLayer
+        Type of FFN layer to use, can be DinoVisionTransformerFFNLayer.MLP,
+        DinoVisionTransformerFFNLayer.SWIGLU, DinoVisionTransformerFFNLayer.SWIGLU_FUSED,
+        or DinoVisionTransformerFFNLayer.IDENTITY. Default is DinoVisionTransformerFFNLayer.MLP.
+    block_chunks : int
+        Split block sequence into block_chunks units for FSDP wrap.
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 0.
+    interpolate_antialias : str
+        Flag to apply anti-aliasing when interpolating positional embeddings.
+    interpolate_offset : float
+        Work-around offset to apply when interpolating positional embeddings.
+    """
+
+    def __init__(
+        self,
+        img_size: int | tuple[int, int, int] = 224,
+        patch_size: int | tuple[int, int, int] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        ffn_bias: bool = True,
+        proj_bias: bool = True,
+        drop_path_rate: float = 0.0,
+        drop_path_uniform: bool = False,
+        init_values: float | None = None,  # for layerscale: None or 0 => no layerscale
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        block_fn: Callable[..., nn.Module] = Block,
+        ffn_layer: DinoVisionTransformerFFNLayer = DinoVisionTransformerFFNLayer.MLP,
+        block_chunks: int = 1,
+        num_register_tokens: int = 0,
+        interpolate_antialias: bool = False,
+        interpolate_offset: float = 0.1,
+    ) -> None:
+        """Inits :class:`DinoVisionTransformer3d`.
+
+        Parameters
+        ----------
+        img_size : int or tuple[int, int, int]
+            Input image size, either a single integer or a tuple of three integers (depth, height, width).
+        patch_size : int or tuple[int, int, int]
+            Patch size, either a single integer or a tuple of three integers (depth, height, width).
+        in_chans : int
+            Number of input channels, default is 3.
+        embed_dim : int
+            Embedding dimension.
+        depth : int
+            Depth of transformer.
+        num_heads : int
+            Number of attention heads.
+        mlp_ratio : int
+            Ratio of mlp hidden dim to embedding dim.
+        qkv_bias : bool
+            Enable bias for qkv if True.
+        proj_bias : bool
+            Enable bias for proj in attn if True.
+        ffn_bias : bool
+            Enable bias for ffn if True.
+        drop_path_rate : float
+            Stochastic depth rate.
+        drop_path_uniform : bool
+            Apply uniform drop rate across blocks.
+        weight_init : str
+            Weight init scheme.
+        init_values : float
+            Layer-scale init values.
+        act_layer : nn.Module
+            MLP activation layer.
+        block_fn : nn.Module
+            Transformer block class.
+        ffn_layer : DinoVisionTransformerFFNLayer
+            Type of FFN layer to use, can be DinoVisionTransformerFFNLayer.MLP,
+            DinoVisionTransformerFFNLayer.SWIGLU, DinoVisionTransformerFFNLayer.SWIGLU_FUSED,
+            or DinoVisionTransformerFFNLayer.IDENTITY. Default is DinoVisionTransformerFFNLayer.MLP.
+        block_chunks : int
+            Split block sequence into block_chunks units for FSDP wrap.
+        num_register_tokens : int
+            Number of extra cls tokens (so-called "registers").
+        interpolate_antialias : str
+            Flag to apply anti-aliasing when interpolating positional embeddings.
+        interpolate_offset : float
+            Work-around offset to apply when interpolating positional embeddings.
+        """
+        super().__init__(
+            dim=DinoVisionTransformerDim.THREE_D,
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            depth=depth,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            ffn_bias=ffn_bias,
+            proj_bias=proj_bias,
+            drop_path_rate=drop_path_rate,
+            drop_path_uniform=drop_path_uniform,
+            init_values=init_values,
+            act_layer=act_layer,
+            block_fn=block_fn,
+            ffn_layer=ffn_layer,
+            block_chunks=block_chunks,
+            num_register_tokens=num_register_tokens,
+            interpolate_antialias=interpolate_antialias,
+            interpolate_offset=interpolate_offset,
+        )
+
+    @override
+    def _interpolate_and_reshape_pos_embed(
+        self,
+        patch_pos_embed: torch.Tensor,
+        patches_resolution: tuple[int, int, int],
+        dim: int,
+        interpolation_kwargs: dict,
+    ) -> torch.Tensor:
+        """Interpolate and reshape 3D patch positional embeddings.
+
+        Parameters
+        ----------
+        patch_pos_embed : torch.Tensor
+            Positional embedding tensor of shape (1, N, C).
+        patches_resolution : tuple of ints
+            Number of patches along each spatial dimension.
+        dim : int
+            Embedding dimension.
+        interpolation_kwargs : dict
+            Arguments passed to `F.interpolate`.
+
+        Returns
+        -------
+        torch.Tensor
+            Reshaped and interpolated tensor of shape (1, D, H, W, C),
+            where D, H, W are the number of patches along depth, height, and width.
+        """
+        patch_pos_embed = patch_pos_embed.reshape(1, *patches_resolution, dim).permute(0, 4, 1, 2, 3)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed,
+            mode="trilinear",
+            antialias=self.interpolate_antialias,
+            **interpolation_kwargs,
+        )
+        return patch_pos_embed.permute(0, 2, 3, 4, 1)
+
+
+def vit_3d_small(
+    patch_size: int | tuple[int, int, int] = (16, 16, 16),
+    num_register_tokens: int = 4,
+    **kwargs,
+) -> DinoVisionTransformer3d:
+    """Builds a small 3d vision transformer with 384-dimensional embeddings, 12 layers, 6 heads, and 4x MLP ratio.
+
+    Parameters
+    ----------
+    patch_size : int or tuple[int, int, int]
+        Patch size, either a single integer or a tuple of three integers (depth, height, width).
+        Default is (16, 16, 16).
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 4.
+    kwargs : dict
+        Additional keyword arguments to pass to the :class:`DinoVisionTransformer3d` constructor.
+
+    Returns
+    -------
+    DinoVisionTransformer3d
+        A small 3d vision transformer.
+    """
+    model = DinoVisionTransformer3d(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_3d_base(
+    patch_size: int | tuple[int, int, int] = (16, 16, 16),
+    num_register_tokens: int = 4,
+    **kwargs,
+) -> DinoVisionTransformer3d:
+    """Builds a base 3d vision transformer with 768-dimensional embeddings, 12 layers, 12 heads, and 4x MLP ratio.
+
+    Parameters
+    ----------
+    patch_size : int or tuple[int, int, int]
+        Patch size, either a single integer or a tuple of three integers (depth, height, width).
+        Default is (16, 16, 16).
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 4.
+    kwargs : dict
+        Additional keyword arguments to pass to the :class:`DinoVisionTransformer3d` constructor.
+
+    Returns
+    -------
+    DinoVisionTransformer3d
+        A base 3d vision transformer.
+    """
+    model = DinoVisionTransformer3d(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_3d_large(
+    patch_size: int | tuple[int, int, int] = (16, 16, 16),
+    num_register_tokens: int = 4,
+    **kwargs,
+) -> DinoVisionTransformer3d:
+    """Builds a large 3d vision transformer with 1024-dimensional embeddings, 24 layers, 16 heads, and 4x MLP ratio.
+
+    Parameters
+    ----------
+    patch_size : int or tuple[int, int, int]
+        Patch size, either a single integer or a tuple of three integers (depth, height, width).
+        Default is (16, 16, 16).
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 4.
+    kwargs : dict
+        Additional keyword arguments to pass to the :class:`DinoVisionTransformer3d` constructor.
+
+    Returns
+    -------
+    DinoVisionTransformer3d
+        A large 3d vision transformer.
+    """
+    model = DinoVisionTransformer3d(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model

vision_transformer_base.py

@@ -0,0 +1,630 @@
+# Copyright 2025 AI for Oncology Research Group. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# References:
+#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from abc import abstractmethod
+from enum import Enum
+import logging
+import math
+from functools import partial
+from typing import Callable, Sequence
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from torch.nn.init import trunc_normal_
+
+from .mlp import Mlp
+from .transformer_block import NestedTensorBlock as Block
+from .patch_embed import PatchEmbed, PatchEmbed3d
+from .swiglu_ffn import SwiGLUFFNFused
+from .helpers import make_2tuple, make_3tuple
+
+logger = logging.getLogger("dinov2")
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+class BlockChunk(nn.ModuleList):
+    """Block chunk for FSDP wrap."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the block chunk.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor.
+        """
+        for b in self:
+            x = b(x)
+        return x
+
+
+class DinoVisionTransformerDim(str, Enum):
+    """Dimension type for DinoVisionTransformer."""
+
+    TWO_D = "2d"
+    THREE_D = "3d"
+
+
+class DinoVisionTransformerFFNLayer(str, Enum):
+    """FFN layer type for DinoVisionTransformer."""
+
+    MLP = "mlp"
+    SWIGLU = "swiglu"
+    SWIGLU_FUSED = "swiglufused"
+    IDENTITY = "identity"
+
+    @classmethod
+    def _missing_(cls, value):
+        if isinstance(value, str):
+            value = value.lower()
+            for member in cls:
+                if member.value == value:
+                    return member
+        raise ValueError(f"{value!r} is not a valid {cls.__name__}")
+
+
+class DinoVisionTransformerBase(nn.Module):
+    """Base class for DinoVisionTransformer, supporting both 2D and 3D vision transformers.
+
+    Parameters
+    ----------
+    dim : DinoVisionTransformerDim
+        Dimension type, either DinoVisionTransformerDim.TWO_D or DinoVisionTransformerDim.THREE_D.
+    img_size : int, tuple[int, int] or tuple[int, int, int]
+        Input image size, either a single integer or a tuple.
+        For 2D, it should be a tuple of two integers (height, width).
+        For 3D, it should be a tuple of three integers (depth, height, width).
+    patch_size : int, tuple[int, int] or tuple[int, int, int]
+        Patch size, either a single integer or a tuple.
+        For 2D, it should be a tuple of two integers (height, width).
+        For 3D, it should be a tuple of three integers (depth, height, width).
+    in_chans : int
+        Number of input channels, default is 3.
+    embed_dim : int
+        Embedding dimension.
+    depth : int
+        Depth of transformer.
+    num_heads : int
+        Number of attention heads.
+    mlp_ratio : int
+        Ratio of mlp hidden dim to embedding dim.
+    qkv_bias : bool
+        Enable bias for qkv if True.
+    proj_bias : bool
+        Enable bias for proj in attn if True.
+    ffn_bias : bool
+        Enable bias for ffn if True.
+    drop_path_rate : float
+        Stochastic depth rate.
+    drop_path_uniform : bool
+        Apply uniform drop rate across blocks.
+    weight_init : str
+        Weight init scheme.
+    init_values : float
+        Layer-scale init values.
+    act_layer : nn.Module
+        MLP activation layer.
+    block_fn : nn.Module
+        Transformer block class.
+    ffn_layer : DinoVisionTransformerFFNLayer
+        Type of FFN layer to use, can be DinoVisionTransformerFFNLayer.MLP,
+        DinoVisionTransformerFFNLayer.SWIGLU, DinoVisionTransformerFFNLayer.SWIGLU_FUSED,
+        or DinoVisionTransformerFFNLayer.IDENTITY. Default is DinoVisionTransformerFFNLayer.MLP.
+    block_chunks : int
+        Split block sequence into block_chunks units for FSDP wrap.
+    num_register_tokens : int
+        Number of extra tokens for the model to deposit information (so-called "registers"). Default is 0.
+    interpolate_antialias : str
+        Flag to apply anti-aliasing when interpolating positional embeddings.
+    interpolate_offset : float
+        Work-around offset to apply when interpolating positional embeddings.
+    """
+
+    def __init__(
+        self,
+        dim: DinoVisionTransformerDim,
+        img_size: int | tuple[int, int] | tuple[int, int, int] = 224,
+        patch_size: int | tuple[int, int] | tuple[int, int, int] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        ffn_bias: bool = True,
+        proj_bias: bool = True,
+        drop_path_rate: float = 0.0,
+        drop_path_uniform: bool = False,
+        init_values: float | None = None,  # for layerscale: None or 0 => no layerscale
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        block_fn: Callable[..., nn.Module] = Block,
+        ffn_layer: DinoVisionTransformerFFNLayer = DinoVisionTransformerFFNLayer.MLP,
+        block_chunks: int = 1,
+        num_register_tokens: int = 0,
+        interpolate_antialias: bool = False,
+        interpolate_offset: float = 0.1,
+    ) -> None:
+        """Inits :class:`DinoVisionTransformerBase`.
+
+        Parameters
+        ----------
+        dim : DinoVisionTransformerDim
+            Dimension type, either DinoVisionTransformerDim.TWO_D or DinoVisionTransformerDim.THREE_D.
+        img_size : int, tuple[int, int] or tuple[int, int, int]
+            Input image size, either a single integer or a tuple.
+            For 2D, it should be a tuple of two integers (height, width).
+            For 3D, it should be a tuple of three integers (depth, height, width).
+        patch_size : int, tuple[int, int] or tuple[int, int, int]
+            Patch size, either a single integer or a tuple.
+            For 2D, it should be a tuple of two integers (height, width).
+            For 3D, it should be a tuple of three integers (depth, height, width).
+        in_chans : int
+            Number of input channels, default is 3.
+        embed_dim : int
+            Embedding dimension.
+        depth : int
+            Depth of transformer.
+        num_heads : int
+            Number of attention heads.
+        mlp_ratio : int
+            Ratio of mlp hidden dim to embedding dim.
+        qkv_bias : bool
+            Enable bias for qkv if True.
+        proj_bias : bool
+            Enable bias for proj in attn if True.
+        ffn_bias : bool
+            Enable bias for ffn if True.
+        drop_path_rate : float
+            Stochastic depth rate.
+        drop_path_uniform : bool
+            Apply uniform drop rate across blocks.
+        weight_init : str
+            Weight init scheme.
+        init_values : float
+            Layer-scale init values.
+        act_layer : nn.Module
+            MLP activation layer.
+        block_fn : nn.Module
+            Transformer block class.
+        ffn_layer : DinoVisionTransformerFFNLayer
+            Type of FFN layer to use, can be DinoVisionTransformerFFNLayer.MLP,
+            DinoVisionTransformerFFNLayer.SWIGLU, DinoVisionTransformerFFNLayer.SWIGLU_FUSED,
+            or DinoVisionTransformerFFNLayer.IDENTITY. Default is DinoVisionTransformerFFNLayer.MLP.
+        block_chunks : int
+            Split block sequence into block_chunks units for FSDP wrap.
+        num_register_tokens : int
+            Number of extra tokens for the model to deposit information (so-called "registers"). Default is 0.
+        interpolate_antialias : str
+            Flag to apply anti-aliasing when interpolating positional embeddings.
+        interpolate_offset : float
+            Work-around offset to apply when interpolating positional embeddings.
+        """
+
+        super().__init__()
+        self.logger = logging.getLogger(type(self).__name__)
+        self.dim = dim
+
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_tokens = 1
+        self.n_blocks = depth
+        self.num_heads = num_heads
+
+        self.patch_size = make_2tuple(patch_size) if dim == DinoVisionTransformerDim.TWO_D else make_3tuple(patch_size)
+        self.img_size = make_2tuple(img_size) if dim == DinoVisionTransformerDim.TWO_D else make_3tuple(img_size)
+
+        if len(self.patch_size) != len(self.img_size):
+            raise ValueError("Patch size and image size must have the same number of dimensions")
+
+        self.num_register_tokens = num_register_tokens
+        self.interpolate_antialias = interpolate_antialias
+        self.interpolate_offset = interpolate_offset
+
+        self.patch_embed = (PatchEmbed if dim == DinoVisionTransformerDim.TWO_D else PatchEmbed3d)(
+            img_size=self.img_size,
+            patch_size=self.patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+        assert num_register_tokens >= 0
+        self.register_tokens = (
+            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
+        )
+
+        if drop_path_uniform is True:
+            dpr = [drop_path_rate] * depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        if ffn_layer == DinoVisionTransformerFFNLayer.MLP:
+            self.logger.info("Using MLP layer as FFN")
+            ffn_layer = Mlp
+        elif (
+            ffn_layer == DinoVisionTransformerFFNLayer.SWIGLU or ffn_layer == DinoVisionTransformerFFNLayer.SWIGLU_FUSED
+        ):
+            self.logger.info("Using SwiGLU layer as FFN")
+            ffn_layer = SwiGLUFFNFused
+        else:  # ffn_layer == DinoVisionTransformerFFNLayer.IDENTITY:
+            self.logger.info("Using Identity layer as FFN")
+            ffn_layer = nn.Identity
+
+        blocks_list = [
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                ffn_layer=ffn_layer,
+                init_values=init_values,
+            )
+            for i in range(depth)
+        ]
+        if block_chunks > 0:
+            self.chunked_blocks = True
+            chunked_blocks = []
+            chunksize = depth // block_chunks
+            for i in range(0, depth, chunksize):
+                # this is to keep the block index consistent if we chunk the block list
+                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+        else:
+            self.chunked_blocks = False
+            self.blocks = nn.ModuleList(blocks_list)
+
+        self.norm = norm_layer(embed_dim)
+        self.head = nn.Identity()
+
+        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        """Initialize weights of the model."""
+        trunc_normal_(self.pos_embed, std=0.02)
+        nn.init.normal_(self.cls_token, std=1e-6)
+        if self.register_tokens is not None:
+            nn.init.normal_(self.register_tokens, std=1e-6)
+        named_apply(init_weights_vit_timm, self)
+
+    def _interpolate_pos_encoding(
+        self, x: torch.Tensor, img_shape: tuple[int, int] | tuple[int, int, int]
+    ) -> torch.Tensor:
+        """Interpolate the positional encoding to match the input image shape.
+
+        This method resizes the positional encoding tensor to match the spatial dimensions of the input tensor.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, N, C) where B is the batch size, N is the number of patches + tokens,
+            and C is the embedding dimension.
+        img_shape : tuple[int, int] | tuple[int, int, int]
+            Spatial dimensions of the input image. For 2D, it should be a tuple of two integers (height, width).
+            For 3D, it should be a tuple of three integers (depth, height, width).
+
+        Returns
+        -------
+        torch.Tensor
+            Interpolated positional encoding tensor of shape (1, N, C), where N is the number of patches + tokens
+        """
+        previous_dtype = x.dtype
+        num_image_patches = x.shape[1] - 1
+
+        N = self.pos_embed.shape[1] - 1
+
+        if num_image_patches == N and all(img_shape[i] == img_shape[i + 1] for i in range(len(img_shape) - 1)):
+            return self.pos_embed
+
+        pos_embed = self.pos_embed.float()
+
+        class_pos_embed = pos_embed[:, 0]
+        patch_pos_embed = pos_embed[:, 1:]
+        dim = x.shape[-1]
+
+        img_shape0 = [img_shape[i] // self.patch_size[i] for i in range(len(img_shape))]
+
+        patches_resolution = self.patch_embed.patches_resolution  # Recover the number of patches in each dimension
+
+        if N != math.prod(patches_resolution):
+            raise ValueError(
+                f"Mismatch: learned pos_embed has {N} tokens, but expected {math.prod(patches_resolution)} patches "
+                f"corresponding to {patches_resolution} resolution."
+            )
+
+        interpolation_kwargs = {}
+        if self.interpolate_offset:
+            scale_factor = [float(s + self.interpolate_offset) / m for (s, m) in zip(img_shape0, patches_resolution)]
+            interpolation_kwargs["scale_factor"] = scale_factor
+        else:
+            # Simply specify an output size instead of a scale factor
+            interpolation_kwargs["size"] = img_shape0
+
+        patch_pos_embed = self._interpolate_and_reshape_pos_embed(
+            patch_pos_embed, patches_resolution, dim, interpolation_kwargs
+        )
+
+        if tuple(img_shape0) != patch_pos_embed.shape[1:-1]:
+            raise ValueError(
+                f"Positional embedding shape mismatch: expected {img_shape0}, got {patch_pos_embed.shape[1:-1]}. "
+                "This may lead to unexpected behavior."
+            )
+
+        patch_pos_embed = patch_pos_embed.view(1, -1, dim)
+
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+    @abstractmethod
+    def _interpolate_and_reshape_pos_embed(
+        self, patch_pos_embed: torch.Tensor, patches_resolution: tuple[int, ...], dim: int, interpolation_kwargs: dict
+    ) -> torch.Tensor:
+        """Subclasses should implement interpolation and reshaping appropriate for 2D or 3D positional embeddings.
+
+        Parameters
+        ----------
+        patch_pos_embed : torch.Tensor
+            Positional embedding tensor of shape (1, N, C).
+        patches_resolution : tuple of ints
+            Number of patches along each spatial dimension.
+        dim : int
+            Embedding dimension.
+        interpolation_kwargs : dict
+            Arguments passed to `F.interpolate`.
+
+        Returns
+        -------
+        torch.Tensor
+            Reshaped and interpolated tensor of shape (1, ..., ..., C).
+        """
+        raise NotImplementedError("Subclasses must implement `_interpolate_and_reshape_pos_embed` method.")
+
+    def _prepare_tokens_with_masks(self, x: torch.Tensor, masks: torch.Tensor | None = None) -> torch.Tensor:
+        """Prepare tokens with masks for the input tensor.
+
+        This method applies patch embedding, adds class tokens, and interpolates positional encodings.
+        If masks are provided, it replaces the corresponding patches with a mask token.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, C, H, W) for 2D or (B, C, D, H, W) for 3D,
+            where B is the batch size, C is the number of channels, and H, W (or D, H, W) are the spatial dimensions.
+        masks : torch.Tensor, optional
+            Optional mask tensor of shape (B, N) where B is the batch size and N is the number of patches.
+            Default is None.
+
+        Returns
+        -------
+        torch.Tensor
+            Prepared tensor of shape (B, N, C) where B is the batch size, N is the number of patches + tokens,
+            and C is the embedding dimension.
+        """
+        x_shape = x.shape[2:]
+        x = self.patch_embed(x)
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        x = x + self._interpolate_pos_encoding(x, x_shape)
+        if self.register_tokens is not None:
+            x = torch.cat(
+                (
+                    x[:, :1],
+                    self.register_tokens.expand(x.shape[0], -1, -1),
+                    x[:, 1:],
+                ),
+                dim=1,
+            )
+
+        return x
+
+    def forward_features_list(
+        self, x_list: list[torch.Tensor], masks_list: list[torch.Tensor]
+    ) -> list[dict[str, torch.Tensor]]:
+        """Forward pass for a list of input tensors with corresponding masks.
+
+        Parameters
+        ----------
+        x_list : list[torch.Tensor]
+            List of input tensors, each of shape (B, C, H, W) for 2D or (B, C, D, H, W) for 3D,
+            where B is the batch size, C is the number of channels, and H, W (or D, H, W) are the spatial dimensions.
+        masks_list : list[torch.Tensor]
+            List of mask tensors, each of shape (B, N) where B is the batch size and N is the number of patches.
+
+        Returns
+        -------
+        list[dict[str, torch.Tensor]]
+            List of dictionaries containing the normalized outputs and masks for each input tensor.
+        """
+        x = [self._prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+        for blk in self.blocks:
+            x = blk(x)
+
+        all_x = x
+        output = []
+        for x, masks in zip(all_x, masks_list):
+            x_norm = self.norm(x)
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm[:, 0],
+                    "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward_features(
+        self,
+        x: torch.Tensor | list[torch.Tensor],
+        masks: torch.Tensor | list[torch.Tensor] = None,
+    ) -> dict[str, torch.Tensor]:
+        """Return features from the input.
+
+        Parameters
+        ----------
+        x : torch.Tensor | list[torch.Tensor]
+            Input tensor or list of input tensors.
+        masks : torch.Tensor | list[torch.Tensor], optional
+            Mask tensor or list of mask tensors.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            Dictionary containing the normalized outputs and masks.
+        """
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self._prepare_tokens_with_masks(x, masks)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x_norm = self.norm(x)
+        return {
+            "x_norm_clstoken": x_norm[:, 0],
+            "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+            "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+            "x_prenorm": x,
+            "masks": masks,
+        }
+
+    def _get_intermediate_layers_not_chunked(self, x: torch.Tensor, n: int | list[int] = 1) -> list[torch.Tensor]:
+        """Get intermediate layers from the transformer blocks."""
+        x = self._prepare_tokens_with_masks(x)
+        # If n is an int, take the n last blocks. If it's a list, take them
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in blocks_to_take:
+                output.append(x)
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def _get_intermediate_layers_chunked(self, x: torch.Tensor, n: int | list[int] = 1) -> list[torch.Tensor]:
+        """Get intermediate layers from the transformer blocks when using chunked blocks."""
+        x = self._prepare_tokens_with_masks(x)
+        output, i, total_block_len = [], 0, len(self.blocks[-1])
+        # If n is an int, take the n last blocks. If it's a list, take them
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for block_chunk in self.blocks:
+            for blk in block_chunk[i:]:  # Passing the nn.Identity()
+                x = blk(x)
+                if i in blocks_to_take:
+                    output.append(x)
+                i += 1
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: int | Sequence = 1,  # Layers or n last layers to take
+        reshape: bool = False,
+        return_class_token: bool = False,
+        norm: bool = True,
+    ) -> tuple[torch.Tensor | tuple[torch.Tensor]]:
+        """Get intermediate layers from the transformer blocks.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor.
+        n : int or Sequence, optional
+            Number of layers or specific layers to take.
+        reshape : bool, optional
+            Whether to reshape the output.
+        return_class_token : bool, optional
+            Whether to return the class token.
+        norm : bool, optional
+            Whether to apply normalization.
+
+        Returns
+        -------
+        tuple[torch.Tensor | tuple[torch.Tensor]]
+            Intermediate layers from the transformer blocks.
+        """
+        if self.chunked_blocks:
+            outputs = self._get_intermediate_layers_chunked(x, n)
+        else:
+            outputs = self._get_intermediate_layers_not_chunked(x, n)
+
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+
+        class_tokens = [out[:, 0] for out in outputs]
+        outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs]
+
+        if reshape:
+            B = x.size(0)
+            spatial_dims = x.shape[2:]
+            outputs = [
+                out.reshape([B] + [s // p for s, p in zip(spatial_dims, self.patch_size)] + [-1])
+                .permute([0] + [x.ndim - 1] + list(range(1, x.ndim - 1)))
+                .contiguous()
+                for out in outputs
+            ]
+
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+
+        return tuple(outputs)
+
+    def forward(self, *args, is_training=False, **kwargs) -> dict[str, torch.Tensor] | torch.Tensor:
+        """Forward pass of :class:`DinoVisionTransformerBase`."""
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return self.head(ret["x_norm_clstoken"])
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = "") -> None:
+    """ViT weight initialization, original timm impl (for reproducibility)"""
+    if isinstance(module, nn.Linear):
+        trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)