"""Utilities for SwitchHead model.

Inspired by: https://github.com/RobertCsordas/moeut
"""

from dataclasses import dataclass

import torch


@dataclass
class AttentionMask:
    """A dataclass to hold attention masks for the SwitchHead model."""

    src_length_mask: torch.Tensor | None
    position_mask: torch.Tensor | None


KVCache = dict[str, torch.Tensor] | None
