shadax 0.2.0__py3-none-any.whl

shadax/__init__.py

@@ -0,0 +1,50 @@
+"""
+SHADAX - Self-supervised Hierarchical Adaptive Hybrid Algorithm.
+
+A production-ready Python implementation of the SHADA algorithm: a four-stage
+hierarchical encoder shared across the image and text modalities, four task
+heads (classification, detection, segmentation, language modeling), two
+self-supervised objectives and a four-phase training pipeline, all behind a
+familiar sklearn-style estimator.
+
+Public API:
+    SHADA: High-level sklearn-style estimator (fit/predict/score).
+    SHADAConfig: Configuration dataclass (the model-shape contract).
+    create_config: Tier-aware configuration factory.
+    SHADANet: The unified encoder + head + SSL ``nn.Module``.
+    HierarchicalEncoder: The shared multi-modal backbone.
+    ModelTier / TaskType / TrainingPhase / Modality: Controlled vocabularies.
+
+Example:
+    >>> from shadax import SHADA
+    >>> model = SHADA(tier="base", num_classes=10)
+    >>> model.fit(X_train, y_train)            # doctest: +SKIP
+    >>> predictions = model.predict(X_test)    # doctest: +SKIP
+"""
+
+__version__ = "0.2.0"
+__author__ = "Omar"
+
+from .config import (
+    Modality,
+    ModelTier,
+    SHADAConfig,
+    TaskType,
+    TrainingPhase,
+    create_config,
+)
+from .encoder import HierarchicalEncoder
+from .model import SHADA
+from .network import SHADANet
+
+__all__ = [
+    "SHADA",
+    "SHADAConfig",
+    "create_config",
+    "ModelTier",
+    "TaskType",
+    "TrainingPhase",
+    "Modality",
+    "HierarchicalEncoder",
+    "SHADANet",
+]

shadax/config.py

@@ -0,0 +1,294 @@
+"""
+SHADA configuration layer.
+
+This module is the *contract* that every other module in :mod:`shadax`
+depends on. It is intentionally free of any heavy dependency (no PyTorch),
+so it can be imported cheaply and reasoned about in isolation.
+
+It defines:
+
+* :class:`ModelTier`, :class:`TaskType`, :class:`TrainingPhase`,
+  :class:`Modality` -- the controlled vocabularies used across the library.
+* :class:`SHADAConfig` -- the single source of truth for the model shape.
+* :data:`_TIER_CONFIGS` -- the per-tier architectural presets.
+* :func:`create_config` -- the factory used by the high level API.
+
+Spatial contract (images)
+--------------------------
+The hierarchical encoder reduces the spatial resolution by a fixed total
+factor of ``32`` (a ``/4`` convolutional stem followed by three ``/2`` stage
+downsamples). Input height and width must therefore be divisible by ``32``;
+:func:`SHADAConfig.validate` and the encoder both enforce this.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, List
+
+__all__ = [
+    "ModelTier",
+    "TaskType",
+    "TrainingPhase",
+    "Modality",
+    "SHADAConfig",
+    "create_config",
+    "STEM_REDUCTION",
+    "STAGE_REDUCTION",
+    "TOTAL_REDUCTION",
+    "NUM_STAGES",
+]
+
+# --------------------------------------------------------------------------- #
+# Architectural constants (the spatial contract).
+# --------------------------------------------------------------------------- #
+NUM_STAGES: int = 4
+STEM_REDUCTION: int = 4            # convolutional stem downsamples H, W by 4.
+STAGE_REDUCTION: int = 2          # each inter-stage downsample halves H, W.
+# stem (/4) + 3 downsamples (/2 each) -> /32 total.
+TOTAL_REDUCTION: int = STEM_REDUCTION * (STAGE_REDUCTION ** (NUM_STAGES - 1))
+
+
+class ModelTier(str, Enum):
+    """Model size tiers, ordered from smallest to largest."""
+
+    NANO = "nano"
+    BASE = "base"
+    LARGE = "large"
+    XL = "xl"
+
+
+class TaskType(str, Enum):
+    """Supported downstream tasks.
+
+    Each value maps to a concrete task head in :mod:`shadax.heads` and to a
+    loss in :mod:`shadax.training`.
+    """
+
+    CLASSIFICATION = "classification"
+    DETECTION = "detection"
+    SEGMENTATION = "segmentation"
+    LANGUAGE_MODEL = "lm"
+
+
+class TrainingPhase(str, Enum):
+    """The four phases of the SHADA training pipeline.
+
+    * ``PRETRAIN``  -- self-supervised pretraining (no labels required).
+    * ``MULTITASK`` -- joint optimisation of the supervised objective and the
+      self-supervised objective.
+    * ``FINETUNE``  -- purely supervised optimisation of the target task.
+    * ``DEPLOY``    -- inference / frozen-weights mode (no optimisation).
+    """
+
+    PRETRAIN = "pretrain"
+    MULTITASK = "multitask"
+    FINETUNE = "finetune"
+    DEPLOY = "deploy"
+
+
+class Modality(str, Enum):
+    """Input modalities understood by the encoder."""
+
+    IMAGE = "image"
+    TEXT = "text"
+
+
+@dataclass
+class SHADAConfig:
+    """Single source of truth for the SHADA architecture.
+
+    Attributes:
+        tier: Model size tier name (``nano``/``base``/``large``/``xl``).
+        encoder_dims: Channel width of each of the four encoder stages.
+        encoder_depths: Number of transformer blocks in each stage.
+        num_heads: Number of attention heads in each stage.
+        mlp_ratio: Hidden expansion ratio of the feed-forward sub-layer.
+        dropout: Dropout probability used throughout the network.
+        in_channels: Number of input image channels (e.g. 3 for RGB).
+        image_size: Reference image size used for defaults; the encoder is
+            resolution-adaptive, so this is only a hint (any H, W divisible by
+            :data:`TOTAL_REDUCTION` works).
+        max_seq_len: Maximum text sequence length supported.
+        vocab_size: Vocabulary size for the text modality / language model.
+        task: Primary :class:`TaskType` value (stored as ``str``).
+        num_classes: Number of output classes. Interpreted per task:
+            classification -> number of labels, detection -> object
+            categories, segmentation -> per-pixel classes. Ignored for ``lm``
+            (which predicts over ``vocab_size``).
+        mask_ratio: Fraction of image patches masked during self-supervised
+            pretraining (masked image modeling).
+        text_mask_ratio: Fraction of text tokens masked during masked language
+            modeling.
+        decoder_dim: Width of the lightweight self-supervised decoder.
+        decoder_depth: Number of blocks in the self-supervised decoder.
+        pad_token_id: Token id treated as padding / ignored by the LM loss.
+    """
+
+    tier: str = "base"
+    encoder_dims: List[int] = field(default_factory=lambda: [128, 256, 512, 1024])
+    encoder_depths: List[int] = field(default_factory=lambda: [3, 4, 6, 3])
+    num_heads: List[int] = field(default_factory=lambda: [4, 8, 16, 32])
+    mlp_ratio: float = 4.0
+    dropout: float = 0.1
+    in_channels: int = 3
+    image_size: int = 224
+    max_seq_len: int = 1024
+    vocab_size: int = 50257
+    task: str = "classification"
+    num_classes: int = 1000
+    mask_ratio: float = 0.75
+    text_mask_ratio: float = 0.15
+    decoder_dim: int = 256
+    decoder_depth: int = 2
+    pad_token_id: int = 0
+
+    def __post_init__(self) -> None:
+        self.validate()
+
+    # ----------------------------------------------------------------- #
+    # Convenience accessors used by the rest of the library.
+    # ----------------------------------------------------------------- #
+    @property
+    def embed_dim(self) -> int:
+        """Width of the first (stem) stage."""
+        return self.encoder_dims[0]
+
+    @property
+    def final_dim(self) -> int:
+        """Width of the last encoder stage (the global feature dimension)."""
+        return self.encoder_dims[-1]
+
+    @property
+    def num_stages(self) -> int:
+        return len(self.encoder_dims)
+
+    @property
+    def task_type(self) -> TaskType:
+        """The :class:`TaskType` enum corresponding to ``self.task``."""
+        return TaskType(self.task)
+
+    def stage_reduction(self, stage: int) -> int:
+        """Total spatial reduction factor at the output of ``stage`` (0-based)."""
+        return STEM_REDUCTION * (STAGE_REDUCTION ** stage)
+
+    def validate(self) -> "SHADAConfig":
+        """Validate internal consistency. Returns ``self`` for chaining.
+
+        Raises:
+            ValueError: if the per-stage lists are inconsistent or any value
+                is out of range.
+        """
+        n = NUM_STAGES
+        for name, seq in (
+            ("encoder_dims", self.encoder_dims),
+            ("encoder_depths", self.encoder_depths),
+            ("num_heads", self.num_heads),
+        ):
+            if len(seq) != n:
+                raise ValueError(
+                    f"{name} must have exactly {n} entries, got {len(seq)}: {seq!r}"
+                )
+        for dim, heads in zip(self.encoder_dims, self.num_heads):
+            if dim % heads != 0:
+                raise ValueError(
+                    f"each encoder dim must be divisible by its head count; "
+                    f"dim={dim} is not divisible by num_heads={heads}"
+                )
+        if self.task not in {t.value for t in TaskType}:
+            raise ValueError(
+                f"task must be one of {[t.value for t in TaskType]}, got {self.task!r}"
+            )
+        if not (0.0 <= self.mask_ratio < 1.0):
+            raise ValueError(f"mask_ratio must be in [0, 1), got {self.mask_ratio}")
+        if not (0.0 <= self.text_mask_ratio < 1.0):
+            raise ValueError(
+                f"text_mask_ratio must be in [0, 1), got {self.text_mask_ratio}"
+            )
+        if self.num_classes < 1:
+            raise ValueError(f"num_classes must be >= 1, got {self.num_classes}")
+        return self
+
+
+# --------------------------------------------------------------------------- #
+# Per-tier presets. Only architectural fields are stored here; task-specific
+# fields (task, num_classes, ...) are supplied by create_config / the user.
+# --------------------------------------------------------------------------- #
+_TIER_CONFIGS: Dict[str, Dict] = {
+    "nano": {
+        "encoder_dims": [64, 128, 256, 512],
+        "encoder_depths": [2, 2, 4, 2],
+        "num_heads": [2, 4, 8, 16],
+        "max_seq_len": 512,
+        "decoder_dim": 128,
+        "decoder_depth": 1,
+    },
+    "base": {
+        "encoder_dims": [128, 256, 512, 1024],
+        "encoder_depths": [3, 4, 6, 3],
+        "num_heads": [4, 8, 16, 32],
+        "max_seq_len": 1024,
+        "decoder_dim": 256,
+        "decoder_depth": 2,
+    },
+    "large": {
+        "encoder_dims": [192, 384, 768, 1536],
+        "encoder_depths": [3, 4, 18, 3],
+        "num_heads": [6, 12, 24, 48],
+        "max_seq_len": 2048,
+        "decoder_dim": 384,
+        "decoder_depth": 2,
+    },
+    "xl": {
+        "encoder_dims": [256, 512, 1024, 2048],
+        "encoder_depths": [3, 4, 24, 3],
+        "num_heads": [8, 16, 32, 64],
+        "max_seq_len": 4096,
+        "decoder_dim": 512,
+        "decoder_depth": 4,
+    },
+}
+
+
+def create_config(
+    tier: str = "base",
+    task: str = "classification",
+    num_classes: int = 1000,
+    **overrides,
+) -> SHADAConfig:
+    """Create a :class:`SHADAConfig` with tier-specific defaults.
+
+    Args:
+        tier: Model tier (``nano``/``base``/``large``/``xl``).
+        task: Task type (``classification``/``detection``/``segmentation``/``lm``).
+        num_classes: Number of output classes (see :class:`SHADAConfig`).
+        **overrides: Any :class:`SHADAConfig` field to override.
+
+    Returns:
+        A validated :class:`SHADAConfig` instance.
+
+    Raises:
+        ValueError: if ``tier`` is unknown or the resulting config is invalid.
+
+    Example:
+        >>> config = create_config("base", task="classification", num_classes=10)
+        >>> config.final_dim
+        1024
+    """
+    if isinstance(tier, ModelTier):
+        tier = tier.value
+    if tier not in _TIER_CONFIGS:
+        raise ValueError(
+            f"tier must be one of {list(_TIER_CONFIGS)}, got {tier!r}"
+        )
+    if isinstance(task, TaskType):
+        task = task.value
+    kwargs = {
+        **_TIER_CONFIGS[tier],
+        "tier": tier,
+        "task": task,
+        "num_classes": num_classes,
+        **overrides,
+    }
+    return SHADAConfig(**kwargs)

shadax/core.py

@@ -0,0 +1,34 @@
+"""
+Backward-compatibility shim for :mod:`shadax.core`.
+
+The original monolithic ``core.py`` has been split into focused modules
+(:mod:`shadax.config`, :mod:`shadax.network`, :mod:`shadax.model`, ...). This
+module preserves the old import surface so that existing code such as::
+
+    from shadax.core import SHADA, SHADAConfig, create_config
+    from shadax.core import ModelTier, TaskType, TrainingPhase
+
+keeps working unchanged. Everything is re-exported from its new home.
+"""
+
+from __future__ import annotations
+
+from shadax.config import (
+    Modality,
+    ModelTier,
+    SHADAConfig,
+    TaskType,
+    TrainingPhase,
+    create_config,
+)
+from shadax.model import SHADA
+
+__all__ = [
+    "SHADA",
+    "SHADAConfig",
+    "create_config",
+    "ModelTier",
+    "TaskType",
+    "TrainingPhase",
+    "Modality",
+]

shadax/encoder.py

@@ -0,0 +1,301 @@
+"""
+SHADA hierarchical multi-modal encoder.
+
+This module defines :class:`HierarchicalEncoder`, the four-stage backbone that
+turns either an image ``(B, C, H, W)`` or a text token sequence ``(B, L)`` into
+a hierarchy of feature maps plus pooled global features.
+
+The encoder is built entirely from the primitives in :mod:`shadax.modules` and
+driven by a :class:`shadax.config.SHADAConfig`. It has two input paths that
+share their per-stage Transformer blocks:
+
+* **Image path** -- a :class:`~shadax.modules.ConvStem` (``/4``), then four
+  stages each consisting of a :class:`~shadax.modules.ConditionalPositionalEncoding`
+  and a list of :class:`~shadax.modules.TransformerBlock` s operating on the
+  flattened tokens, with a :class:`~shadax.modules.StageDownsample` (``/2``)
+  between consecutive stages. Total spatial reduction is
+  :data:`~shadax.config.TOTAL_REDUCTION` (= 32).
+* **Text path** -- a :class:`~shadax.modules.TextEmbedding`, then the same four
+  stages of Transformer blocks (text tokens are already a sequence), with a
+  length-preserving :class:`~shadax.modules.TextStageProject` between stages.
+
+The Transformer blocks are shared across modalities: both paths ultimately call
+them on token sequences ``(B, N, dim)``.
+
+The dictionary returned by :meth:`HierarchicalEncoder.forward` is a hard
+contract consumed by :mod:`shadax.heads` and :mod:`shadax.ssl`; its keys and
+shapes must not change.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from shadax.config import (
+    NUM_STAGES,
+    STEM_REDUCTION,
+    TOTAL_REDUCTION,
+    Modality,
+    SHADAConfig,
+)
+from shadax.modules import (
+    ConditionalPositionalEncoding,
+    ConvStem,
+    StageDownsample,
+    TextEmbedding,
+    TextStageProject,
+    TransformerBlock,
+)
+
+
+__all__ = ["HierarchicalEncoder"]
+
+
+class _Stage(nn.Module):
+    """A single encoder stage: positional encoding plus Transformer blocks.
+
+    The conditional positional encoding is only used by the image path (it
+    operates on spatial maps); the Transformer block list is shared between the
+    image and text paths, both of which feed it token sequences.
+
+    Args:
+        dim: Channel / token width of this stage.
+        depth: Number of Transformer blocks in this stage.
+        num_heads: Number of attention heads per block.
+        mlp_ratio: Hidden expansion ratio of each block's feed-forward layer.
+        dropout: Dropout probability used throughout the stage.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        dropout: float,
+    ) -> None:
+        super().__init__()
+        self.cpe = ConditionalPositionalEncoding(dim)
+        self.blocks = nn.ModuleList(
+            [
+                TransformerBlock(dim, num_heads, mlp_ratio=mlp_ratio, dropout=dropout)
+                for _ in range(depth)
+            ]
+        )
+
+    def run_blocks(
+        self, tokens: torch.Tensor, attn_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """Run every Transformer block in this stage over a token sequence.
+
+        Args:
+            tokens: Token sequence of shape ``(B, N, dim)``.
+            attn_mask: Optional additive float attention mask ``(N, N)``.
+
+        Returns:
+            Token sequence of shape ``(B, N, dim)``.
+        """
+        for block in self.blocks:
+            tokens = block(tokens, attn_mask=attn_mask)
+        return tokens
+
+
+class HierarchicalEncoder(nn.Module):
+    """Four-stage hierarchical encoder shared across image and text modalities.
+
+    Built from a :class:`~shadax.config.SHADAConfig`, the encoder exposes a
+    single :meth:`forward` that branches on ``modality``. Both branches produce
+    four feature maps of increasing channel width (``config.encoder_dims``),
+    plus a final token sequence and pooled global features.
+
+    For the image modality, ``H`` and ``W`` must each be divisible by
+    :data:`~shadax.config.TOTAL_REDUCTION` (= 32).
+
+    Args:
+        config: The validated SHADA configuration describing the model shape.
+    """
+
+    def __init__(self, config: SHADAConfig) -> None:
+        super().__init__()
+        self.config = config
+        dims: List[int] = config.encoder_dims
+        depths: List[int] = config.encoder_depths
+        heads: List[int] = config.num_heads
+        mlp_ratio: float = config.mlp_ratio
+        dropout: float = config.dropout
+
+        # Image entry point.
+        self.stem = ConvStem(config.in_channels, dims[0], dropout=dropout)
+
+        # Text entry point.
+        self.text_embed = TextEmbedding(
+            vocab_size=config.vocab_size,
+            embed_dim=dims[0],
+            max_seq_len=config.max_seq_len,
+            dropout=dropout,
+            pad_token_id=config.pad_token_id,
+        )
+
+        # Per-stage modules (Transformer blocks shared across modalities).
+        self.stages = nn.ModuleList(
+            [
+                _Stage(dims[i], depths[i], heads[i], mlp_ratio, dropout)
+                for i in range(NUM_STAGES)
+            ]
+        )
+
+        # Inter-stage reductions: image spatial downsamples and text channel
+        # projections (3 each, one between every pair of consecutive stages).
+        self.downsamples = nn.ModuleList(
+            [StageDownsample(dims[i], dims[i + 1]) for i in range(NUM_STAGES - 1)]
+        )
+        self.text_projects = nn.ModuleList(
+            [TextStageProject(dims[i], dims[i + 1]) for i in range(NUM_STAGES - 1)]
+        )
+
+    # ------------------------------------------------------------------ #
+    # Forward dispatch.
+    # ------------------------------------------------------------------ #
+    def forward(
+        self,
+        x: torch.Tensor,
+        modality: Union[str, Modality] = "image",
+        causal: bool = False,
+    ) -> Dict[str, Any]:
+        """Encode an input into a hierarchy of features.
+
+        Args:
+            x: Either an image ``(B, in_channels, H, W)`` (image modality) or a
+                long tensor of token ids ``(B, L)`` (text modality).
+            modality: ``"image"``/``"text"`` (or the corresponding
+                :class:`~shadax.config.Modality` enum value).
+            causal: Text only. If ``True``, a causal attention mask is applied
+                so each position may attend only to itself and earlier
+                positions.
+
+        Returns:
+            A dict with exactly the keys ``"feature_maps"`` (list of 4
+            tensors), ``"tokens"`` ``(B, N_last, dims[-1])``,
+            ``"global_features"`` ``(B, dims[-1])``, ``"hw"`` (an ``(Hs, Ws)``
+            int tuple for images or ``None`` for text) and ``"modality"`` (the
+            resolved modality string).
+
+        Raises:
+            ValueError: for the image modality if ``H`` or ``W`` is not
+                divisible by :data:`~shadax.config.TOTAL_REDUCTION`, or if
+                ``modality`` is not a recognised value.
+        """
+        modality = Modality(modality)
+        if modality is Modality.IMAGE:
+            return self._forward_image(x)
+        if modality is Modality.TEXT:
+            return self._forward_text(x, causal=causal)
+        raise ValueError(f"unsupported modality: {modality!r}")
+
+    def forward_features(
+        self,
+        x: torch.Tensor,
+        modality: Union[str, Modality] = "image",
+        causal: bool = False,
+    ) -> List[torch.Tensor]:
+        """Return only the list of four hierarchical feature maps.
+
+        Convenience wrapper around :meth:`forward` that discards everything
+        except the ``"feature_maps"`` entry.
+
+        Args:
+            x: Input image or token ids (see :meth:`forward`).
+            modality: Input modality (see :meth:`forward`).
+            causal: Whether to apply a causal mask (text only).
+
+        Returns:
+            The list of four feature-map tensors.
+        """
+        return self.forward(x, modality=modality, causal=causal)["feature_maps"]
+
+    # ------------------------------------------------------------------ #
+    # Modality-specific paths.
+    # ------------------------------------------------------------------ #
+    def _forward_image(self, x: torch.Tensor) -> Dict[str, Any]:
+        """Image path of :meth:`forward` (see its docstring)."""
+        _, _, height, width = x.shape
+        if height % TOTAL_REDUCTION != 0 or width % TOTAL_REDUCTION != 0:
+            raise ValueError(
+                f"image height and width must each be divisible by "
+                f"{TOTAL_REDUCTION}, got H={height}, W={width}"
+            )
+
+        m = self.stem(x)  # (B, dims[0], H/4, W/4)
+        feature_maps: List[torch.Tensor] = []
+
+        for i, stage in enumerate(self.stages):
+            m = stage.cpe(m)
+            b, c, h, w = m.shape
+            tokens = m.flatten(2).transpose(1, 2)  # (B, H*W, C)
+            tokens = stage.run_blocks(tokens)
+            m = tokens.transpose(1, 2).reshape(b, c, h, w)  # (B, C, H, W)
+            feature_maps.append(m)
+            if i < NUM_STAGES - 1:
+                m = self.downsamples[i](m)
+
+        last = feature_maps[-1]
+        b, c, h, w = last.shape
+        tokens = last.flatten(2).transpose(1, 2)  # (B, N_last, dims[-1])
+        global_features = tokens.mean(dim=1)  # (B, dims[-1])
+
+        return {
+            "feature_maps": feature_maps,
+            "tokens": tokens,
+            "global_features": global_features,
+            "hw": (h, w),
+            "modality": Modality.IMAGE.value,
+        }
+
+    def _forward_text(self, x: torch.Tensor, causal: bool = False) -> Dict[str, Any]:
+        """Text path of :meth:`forward` (see its docstring)."""
+        t = self.text_embed(x)  # (B, L, dims[0])
+        seq_len = t.shape[1]
+
+        attn_mask: Optional[torch.Tensor] = None
+        if causal:
+            attn_mask = self._causal_mask(seq_len, device=t.device, dtype=t.dtype)
+
+        feature_maps: List[torch.Tensor] = []
+        for i, stage in enumerate(self.stages):
+            t = stage.run_blocks(t, attn_mask=attn_mask)  # (B, L, dims[i])
+            feature_maps.append(t)
+            if i < NUM_STAGES - 1:
+                t = self.text_projects[i](t)  # (B, L, dims[i+1])
+
+        tokens = t  # (B, L, dims[-1])
+        global_features = tokens.mean(dim=1)  # (B, dims[-1])
+
+        return {
+            "feature_maps": feature_maps,
+            "tokens": tokens,
+            "global_features": global_features,
+            "hw": None,
+            "modality": Modality.TEXT.value,
+        }
+
+    @staticmethod
+    def _causal_mask(
+        seq_len: int, device: torch.device, dtype: torch.dtype
+    ) -> torch.Tensor:
+        """Build an additive causal attention mask.
+
+        Args:
+            seq_len: Sequence length ``L``.
+            device: Device of the produced mask.
+            dtype: Floating dtype of the produced mask.
+
+        Returns:
+            A ``(L, L)`` float tensor that is ``0`` on and below the diagonal
+            and ``-inf`` above it, suitable for passing as ``attn_mask`` to
+            :class:`~shadax.modules.TransformerBlock`.
+        """
+        mask = torch.full((seq_len, seq_len), float("-inf"), device=device, dtype=dtype)
+        return torch.triu(mask, diagonal=1)