initial commit with working prototype

.dockerignore

@@ -0,0 +1,67 @@
+# =========================================
+# 1. Python cache & build artifacts
+# =========================================
+__pycache__/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+*.pyc
+*.pyo
+*.pyd
+*.so
+build/
+dist/
+*.egg-info/
+
+# =========================================
+# 2. Virtual environments
+# =========================================
+env/
+.venv/
+venv/
+condaenv.*.requirements.txt
+
+# =========================================
+# 3. Jupyter notebooks
+# =========================================
+.ipynb_checkpoints/
+*.nbconvert.ipynb
+
+# =========================================
+# 4. Logs & temp files
+# =========================================
+*.log
+*.tmp
+*.swp
+*.swo
+*.bak
+.DS_Store
+Thumbs.db
+
+# =========================================
+# 5. IDE / Editor configs
+# =========================================
+.vscode/
+.idea/
+*.sublime-*
+*.iml
+
+# =========================================
+# 6. Machine Learning artifacts
+# =========================================
+data/
+outputs/
+runs/
+artifacts/
+*.pt
+*.pth
+*.ckpt
+*.h5
+*.onnx
+
+# =========================================
+# 7. Version control
+# =========================================
+.git
+.gitignore
+.pre-commit-config.yaml

.gitignore

@@ -0,0 +1,61 @@
+# =========================================
+# Python cache & build artifacts
+# =========================================
+__pycache__/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+*.pyc
+*.pyo
+*.pyd
+*.so
+build/
+dist/
+*.egg-info/
+
+
+# =========================================
+# Virtual / Conda environments
+# =========================================
+env/
+.venv/
+venv/
+condaenv.*.requirements.txt
+
+# =========================================
+# Jupyter / Notebooks
+# =========================================
+.ipynb_checkpoints/
+*.nbconvert.ipynb
+
+# =========================================
+# Logs / Temporary files
+# =========================================
+*.log
+*.tmp
+*.swp
+*.swo
+*.bak
+.DS_Store
+Thumbs.db
+
+# =========================================
+# IDE / Editor settings
+# =========================================
+.vscode/
+.idea/
+*.sublime-*
+*.iml
+
+# =========================================
+# Machine Learning artifacts
+# =========================================
+data/
+outputs/
+runs/
+artifacts/
+*.pt
+*.pth
+*.ckpt
+*.h5
+*.onnx

.pre-commit-config.yaml

@@ -0,0 +1,17 @@
+repos:
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.12.12
+  hooks:
+    - id: ruff           # lints
+    - id: ruff-format    # formats (ruff's formatter is black-compatible)
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v6.0.0
+  hooks:
+    - id: end-of-file-fixer
+    - id: trailing-whitespace
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v1.17.1
+  hooks:
+    - id: mypy
+      args: ["--config-file=pyproject.toml"]
+      additional_dependencies: [pydantic>=2]

Dockerfile.dev

@@ -0,0 +1,38 @@
+# syntax=docker/dockerfile:1
+FROM python:3.12-slim
+# FROM = choose a base OS + language runtime (Debian slim with Python 3.12)
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_NO_CACHE_DIR=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git build-essential && \
+    rm -rf /var/lib/apt/lists/*
+# RUN = execute at build time
+# apt-get update            -> refresh package index
+# apt-get install           -> install tools you might need (git, compilers)
+# --no-install-recommends   -> fewer extra packages
+# rm -rf ...                -> clean apt cache to shrink the layer
+
+WORKDIR /app
+# WORKDIR = default directory for the following Dockerfile steps and 'docker run' processes
+
+COPY requirements.txt pyproject.toml ./
+# COPY = bring files from host into the image
+# Only copying dependency metadata first lets Docker cache layer if code changes later
+
+RUN pip install -U pip && pip install -r requirements.txt
+# pip install -U pip    -> upgrade pip in the image
+# pip install -r ...    -> install your Python runtime dependencies
+
+COPY . .
+# Now copy your actual code (done AFTER deps for better caching)
+
+RUN pip install -e .[dev]
+# -e . = editable install (your src is importable as a package without rebuilding)
+
+CMD ["pytest"]
+# CMD = default command when a container starts from this image
+# Here: run unit tests quietly (-q) in tests/units

Makefile

@@ -0,0 +1,42 @@
+.PHONY: help conda-setup precommit lint fmt type test clean docker-build docker-run
+
+help:
+	@echo "Common Makefile commands:"
+	@echo "  conda-setup     - Create and activate conda env, install all dependencies, setup pre-commit"
+	@echo "  precommit       - Run all pre-commit hooks (lint, fmt, type)"
+	@echo "  lint            - Run ruff linter"
+	@echo "  fmt             - Run ruff formatter and black"
+	@echo "  type            - Run mypy type checks"
+	@echo "  test            - Run all tests with pytest and coverage"
+	@echo "  clean           - Remove build, cache, and venv artifacts"
+	@echo "  docker-build    - Build the Docker dev image"
+	@echo "  docker-run      - Run the Docker dev container (interactive)"
+
+conda-setup:
+	conda env create -n transformer -f environment.yml|| conda env update -n transformer -f environment.yml
+	conda run -n transformer pip install -e .[dev]
+	conda run -n transformer pre-commit install
+
+precommit:
+	conda run -n transformer pre-commit run --all-files
+
+lint:
+	conda run -n transformer ruff check .
+
+fmt:
+	conda run -n transformer ruff format .
+	conda run -n transformer black src tests
+
+type:
+	conda run -n transformer mypy src
+
+test:
+	conda run -n transformer pytest
+clean:
+	rm -rf .venv .pytest_cache .mypy_cache dist build *.egg-info
+
+docker-build:
+	docker compose build dev --no-cache
+
+docker-run:
+	docker compose run --rm dev bash

README.md

@@ -0,0 +1,8 @@
+# transformer
+
+Minimal encoder–decoder Transformer.
+
+## Quickstart
+```bash
+make conda-setup
+make test

architecture.txt

@@ -0,0 +1,28 @@
+Transformer:
+    embedding:
+        input_embedding (same for encoder and decoder)
+        positional_embedding
+        dropout
+    encoder
+        encoder_layer *N
+            multihead_self_attention
+            dropout+norm
+            feedforward:
+                linear
+                relu
+                linear
+            dropout+norm
+    decoder
+        decoder_layer*N
+            masked_multihead_self_attention
+            dropout+norm
+            multihead_cross_attention
+            dropout+norm
+            feedforward:
+                linear
+                relu
+                linear
+            dropout+norm
+    lm_head:
+        linear (same parameters as the embedding)
+        softmax

configs/infer_mode.yaml

@@ -0,0 +1,10 @@
+defaults:
+  - model: test_model
+  - runtime: runtime
+  - _self_
+input_text: ""
+max_new_tokens: 64
+temperature: 1.0
+top_k: 4
+top_p: 0.7
+do_sample: True

configs/model/test_model.yaml

@@ -0,0 +1,10 @@
+d_model: 256
+num_heads: 4
+num_layers: 4
+d_ff: 1024
+dropout_rate: 0.1
+vocab_size: 32100
+max_seq_len: 128
+pad_id: 0
+bos_id: 1
+eos_id: 2

configs/runtime/runtime.yaml

@@ -0,0 +1,3 @@
+tokenizer_name: t5-small
+weights_path: artifacts/checkpoints/initial_model_weights.pth
+artifact_dir: outputs/${now:%Y-%m-%d}/${now:%H-%M-%S}

configs/train_mode.yaml

@@ -0,0 +1,9 @@
+seed: 42
+batch_size: 16
+epochs: 3
+lr: 3e-4
+weight_decay: 0.01
+grad_clip: 1.0
+precision: "fp32"   # one of: fp32, fp16, bf16
+checkpoint_dir: "artifacts/checkpoints"
+log_dir: "artifacts/mlruns"

docker-compose.yml

@@ -0,0 +1,14 @@
+services:
+  dev:
+    build:
+      context: .
+      dockerfile: Dockerfile.dev
+    image: transformer:dev
+    container_name: transformer_dev
+    working_dir: /app
+    volumes:
+      - ./artifacts:/app/artifacts
+    stdin_open: true
+    tty: true
+    gpus: all
+    command: bash

environment.yml

@@ -0,0 +1,12 @@
+name: transformer
+channels:
+  - defaults
+dependencies:
+  - python=3.12
+  - pip
+  - pip:
+      - --extra-index-url https://download.pytorch.org/whl/cu129
+      - torch==2.8.0+cu129
+      - transformers
+      - omegaconf
+      - hydra-core

pyproject.toml

@@ -0,0 +1,58 @@
+[project]
+name = "transformer"
+version = "0.1.0"
+description = "A transformer implementation"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+  "torch==2.8",
+  "transformers",
+  "omegaconf",
+  "hydra-core",
+]
+
+[project.optional-dependencies]
+dev = [
+  "pytest",
+  "pytest-cov",
+  "ruff",
+  "black",
+  "mypy",
+  "pre-commit",
+]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+minversion = "7.0"
+addopts = "-ra -q"
+testpaths = ["tests"]
+
+[tool.black]
+line-length = 100
+target-version = ["py310"]
+
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+select = ["E","F","I","B","UP"]
+ignore = ["E501"]
+
+[tool.mypy]
+python_version = "3.11"
+strict = false
+ignore_missing_imports = true
+
+# Make mypy understand the src/ layout
+files = ["src"]
+mypy_path = ["src"]
+explicit_package_bases = true

scripts/inference.py

@@ -0,0 +1,53 @@
+import hydra
+import torch
+from hydra.utils import to_absolute_path
+from omegaconf import DictConfig, OmegaConf
+from transformers import AutoTokenizer
+
+from transformer import (
+    BasicEncDecCfg,
+    BasicEncoderDecoderTransformer,
+    InferAppCfg,
+)
+
+
+@hydra.main(config_path="../configs", config_name="infer_mode", version_base=None)
+def main(cfg: DictConfig):
+    scfg_temp = OmegaConf.merge(OmegaConf.structured(InferAppCfg), cfg)
+    scfg: InferAppCfg = OmegaConf.to_object(scfg_temp)
+    model_cfg = BasicEncDecCfg(**vars(scfg.model))
+
+    transformer = BasicEncoderDecoderTransformer(model_cfg)
+    tokenizer = AutoTokenizer.from_pretrained(
+        scfg.runtime.tokenizer_name,
+        use_fast=True,
+    )
+
+    if scfg.runtime.weights_path:
+        ckpt = torch.load(to_absolute_path(scfg.runtime.weights_path), map_location="cpu")
+        state = ckpt.get("model_state_dict", ckpt)
+        transformer.load_state_dict(state, strict=False)
+    transformer.eval()
+
+    text = scfg.input_text
+    if not text:
+        raise SystemExit('Pass text like: input_text="hello world"')
+
+    encoded = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
+    src_ids = encoded["input_ids"]
+    src_padd_mask = (encoded["attention_mask"] == 0).unsqueeze(1).unsqueeze(2)
+    tgt_ids = transformer.generate(
+        src_ids,
+        src_padd_mask,
+        max_new_tokens=scfg.max_new_tokens,
+        temperature=scfg.temperature,
+        top_k=scfg.top_k,
+        top_p=scfg.top_p,
+        do_sample=scfg.do_sample,
+    )
+    output = tokenizer.batch_decode(tgt_ids, skip_special_tokens=True)
+    print(output)
+
+
+if __name__ == "__main__":
+    main()

src/transformer/__init__.py

@@ -0,0 +1,4 @@
+from .configs import BasicEncDecCfg, InferAppCfg, TrainAppCfg
+from .transformer import BasicEncoderDecoderTransformer
+
+__all__ = ["BasicEncoderDecoderTransformer", "BasicEncDecCfg", "TrainAppCfg", "InferAppCfg"]

src/transformer/configs.py

@@ -0,0 +1,78 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class BasicEncDecCfg:
+    vocab_size: int
+    d_model: int
+    max_seq_len: int
+    num_heads: int
+    d_ff: int
+    num_layers: int
+    dropout_rate: float = 0.1
+    pad_id: int = 0
+    bos_id: int = 2
+    eos_id: int = 1
+
+
+@dataclass
+class DataCfg:
+    train_path: str
+    val_path: str
+    num_workers: int = 4
+    batch_size: int = 32
+    max_length: int = 512
+
+
+@dataclass
+class OptimCfg:
+    lr: float = 3e-4
+    betas: tuple[float, float] = (0.9, 0.95)
+    weight_decay: float = 0.1
+    eps: float = 1e-8
+
+
+@dataclass
+class SchedCfg:
+    type: str = "cosine"
+    warmup_steps: int = 2000
+    total_steps: int = 200000
+    min_lr: float = 1e-5
+
+
+@dataclass
+class TrainerCfg:
+    epochs: int = 10
+    gradient_accumulation: int = 1
+    eval_interval: int = 1000
+    save_dir: str = "outputs/run"
+
+
+@dataclass
+class RuntimeCfg:
+    tokenizer_name: str = "t5-small"
+    weights_path: str = ""
+    artifact_dir: str = "outputs/run"
+
+
+@dataclass
+class TrainAppCfg:
+    seed: int
+    model: BasicEncDecCfg
+    data: DataCfg
+    optim: OptimCfg
+    scheduler: SchedCfg
+    trainer: TrainerCfg
+    runtime: RuntimeCfg
+
+
+@dataclass
+class InferAppCfg:
+    model: BasicEncDecCfg
+    runtime: RuntimeCfg
+    input_text: str = ""
+    max_new_tokens: int = 64
+    temperature: float = 1.0
+    top_k: int = 0
+    top_p: float = 1.0
+    do_sample: bool = False

src/transformer/modules/__init__.py

@@ -0,0 +1,6 @@
+from transformer.modules.decoder import TransformerDecoder
+from transformer.modules.embedding import InputEmbedding
+from transformer.modules.encoder import TransformerEncoder
+from transformer.modules.lm_head import LMHead
+
+__all__ = ["TransformerEncoder", "TransformerDecoder", "InputEmbedding", "LMHead"]

src/transformer/modules/attention.py

@@ -0,0 +1,105 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from transformer.utils import calculate_attention, join_heads, split_heads
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-head attention: linear projections -> split heads -> scaled dot-product
+    attention (via utils) -> merge heads -> output projection (+ dropout).
+
+    Args:
+        d_model (int): Model dimension (>0). Must be divisible by num_heads.
+        num_heads (int): Number of attention heads (>0).
+        dropout_rate (float): Dropout probability in (0,1).
+
+    Inputs:
+        query, key, value: (B, S, D) with D == d_model
+        mask (optional): Tensor broadcastable to (B, H, S_q, S_k), either boolean
+                         (True = masked) or additive float mask.
+
+    Returns:
+        Tensor: (B, S_q, D)
+    """
+
+    def __init__(self, d_model: int, num_heads: int, dropout_rate: float):
+        super().__init__()
+        # ---- type checks
+        if not isinstance(num_heads, int):
+            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        # ---- value checks
+        if num_heads <= 0:
+            raise ValueError(f"num_heads must be strictly greater than 0, got {num_heads}")
+        if d_model <= 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
+        if d_model % num_heads != 0:
+            raise ValueError("d_model must be divisible by num_heads")
+        if not (0 <= dropout_rate <= 1):
+            raise ValueError(f"dropout_rate must be between 0 and 1 included, got {dropout_rate}")
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_head = d_model // num_heads
+        self.dropout_rate = dropout_rate
+
+        self.query_linear = nn.Linear(d_model, d_model, bias=False)
+        self.key_linear = nn.Linear(d_model, d_model, bias=False)
+        self.value_linear = nn.Linear(d_model, d_model, bias=False)
+        self.output_linear = nn.Linear(d_model, d_model, bias=True)
+        self.dropout = nn.Dropout(dropout_rate)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        mask: Tensor | None,
+    ) -> Tensor:
+        # ---- basic type checks
+        if not isinstance(query, torch.Tensor):
+            raise TypeError(f"query must be a torch.Tensor, got {type(query)}")
+        if not isinstance(key, torch.Tensor):
+            raise TypeError(f"key must be a torch.Tensor, got {type(key)}")
+        if not isinstance(value, torch.Tensor):
+            raise TypeError(f"value must be a torch.Tensor, got {type(value)}")
+        if mask is not None and not isinstance(mask, torch.Tensor):
+            raise TypeError(f"mask must be a torch.Tensor or None, got {type(mask)}")
+
+        # ---- shape checks for q/k/v (3D [B,S,D] and D == d_model)
+        if query.dim() != 3 or key.dim() != 3 or value.dim() != 3:
+            raise ValueError(
+                "query/key/value must be 3D tensors of shape (B, S, D); "
+                f"got q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}"
+            )
+        Bq, Sq, Dq = query.shape
+        Bk, Sk, Dk = key.shape
+        Bv, Sv, Dv = value.shape
+
+        if not (Dq == Dk == Dv == self.d_model):
+            raise ValueError(
+                f"Last dimension must equal d_model={self.d_model}; got Dq={Dq}, Dk={Dk}, Dv={Dv}"
+            )
+        if not (Bq == Bk == Bv):
+            raise ValueError(f"Batch size mismatch: q={Bq}, k={Bk}, v={Bv}")
+        if Sk != Sv:
+            raise ValueError(f"Key/Value seq length mismatch: Sk={Sk} vs Sv={Sv}")
+
+        # ---- project and split into heads -> (B, H, S, Dh)
+        q = split_heads(self.query_linear(query), self.num_heads)
+        k = split_heads(self.key_linear(key), self.num_heads)
+        v = split_heads(self.value_linear(value), self.num_heads)
+
+        # ---- attention (utils handles mask broadcasting/device and numeric stability)
+        p = self.dropout_rate if self.training else 0.0
+        attn = calculate_attention(q, k, v, mask, attn_dropout_p=p)  # (B, H, Sq, Dh)
+
+        # ---- merge heads, project out, dropout
+        out = join_heads(attn)  # (B, Sq, D)
+        return self.dropout(self.output_linear(out))

src/transformer/modules/decoder.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from transformer.modules.attention import MultiHeadAttention
+from transformer.modules.feedforward import FeedForwardLayer
+from transformer.utils import combine_masks
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float):
+        super().__init__()
+
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if not isinstance(num_heads, int):
+            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
+        if not isinstance(d_ff, int):
+            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        if not d_model > 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
+        if not num_heads > 0:
+            raise ValueError(f"num_heads must be strictly greater than 0, got {num_heads}")
+        if not d_ff > 0:
+            raise ValueError(f"d_ff must be strictly greater than 0, got {d_ff}")
+        if not (0.0 <= dropout_rate < 1.0):
+            raise ValueError(f"dropout_rate must be between 0 and 1 excluded, got {dropout_rate}")
+
+        self.self_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
+        self.feed_forward = FeedForwardLayer(d_model, d_ff)
+        self.cross_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout2 = nn.Dropout(dropout_rate)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout3 = nn.Dropout(dropout_rate)
+
+    def forward(
+        self,
+        x: Tensor,
+        y: Tensor,
+        src_padding_mask: Tensor | None,
+        tgt_padding_mask: Tensor | None,
+        tgt_causal_mask: Tensor | None,
+    ) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
+        if not isinstance(y, torch.Tensor):
+            raise TypeError(f"y must be a torch.Tensor, got {type(y)}")
+        if not (x.dim() == 3):
+            raise ValueError(
+                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+            )
+        if not (y.dim() == 3):
+            raise ValueError(
+                f"y must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(y.shape)}"
+            )
+        if not (x.shape[0] == y.shape[0] and x.shape[-1] == y.shape[-1]):
+            raise ValueError(
+                "Batch size or d_model mismatch between encoder memory and decoder input"
+            )
+
+        tgt_mask = combine_masks(tgt_padding_mask, tgt_causal_mask)
+        y = self.norm1(y + self.dropout1(self.self_attention_layer(y, y, y, tgt_mask)))
+        y = self.norm2(y + self.dropout2(self.cross_attention_layer(y, x, x, src_padding_mask)))
+        y = self.norm3(y + self.dropout3(self.feed_forward(y)))
+        return y
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(
+        self, d_model: int, num_heads: int, d_ff: int, num_layers: int, dropout_rate: float
+    ):
+        super().__init__()
+
+        if not isinstance(num_layers, int):
+            raise TypeError(f"num_layers must be an int, got {type(num_layers)}")
+        if not num_layers > 0:
+            raise ValueError(f"num_layers must be strictly greater than 0, got {num_layers}")
+
+        self.layers = nn.ModuleList(
+            [DecoderLayer(d_model, num_heads, d_ff, dropout_rate) for _ in range(num_layers)]
+        )
+
+    def forward(
+        self,
+        x: Tensor,
+        y: Tensor,
+        src_padding_mask: Tensor | None,
+        tgt_padding_mask: Tensor | None,
+        tgt_causal_mask: Tensor | None,
+    ) -> Tensor:
+        for layer in self.layers:
+            y = layer(x, y, src_padding_mask, tgt_padding_mask, tgt_causal_mask)
+        return y

src/transformer/modules/embedding.py

@@ -0,0 +1,121 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from transformer.utils import sinusoidal_positional_encoding
+
+
+class InputEmbedding(nn.Module):
+    """
+    Token + positional embedding module.
+
+    Args:
+        vocab_size (int): Size of vocabulary (>0).
+        d_model (int): Embedding dimension (>0).
+        sequence_length (int): Maximum sequence length (≥0).
+        pad_id (int): Padding token ID (0 <= pad_id < vocab_size).
+
+    Input:
+        x (LongTensor): shape (batch_size, seq_len), token IDs.
+
+    Output:
+        Tensor: shape (batch_size, seq_len, d_model), embeddings.
+
+    Notes:
+        - Zero-length sequences are allowed (returns [B, 0, D]).
+        - Positional encodings are sinusoidal and added to token embeddings.
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int,
+        sequence_length: int,
+        pad_id: int,
+        dropout_rate: float = 0.0,
+    ):
+        super().__init__()
+        # --- type checks
+        if not isinstance(vocab_size, int):
+            raise TypeError(f"vocab_size must be int, got {type(vocab_size)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be int, got {type(d_model)}")
+        if not isinstance(sequence_length, int):
+            raise TypeError(f"sequence_length must be int, got {type(sequence_length)}")
+        if not isinstance(pad_id, int):
+            raise TypeError(f"pad_id must be int, got {type(pad_id)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        # --- value checks
+        if vocab_size <= 0:
+            raise ValueError(f"vocab_size must be > 0, got {vocab_size}")
+        if d_model <= 0:
+            raise ValueError(f"d_model must be > 0, got {d_model}")
+        if sequence_length < 0:
+            raise ValueError(f"sequence_length must be >= 0, got {sequence_length}")
+        if not (0 <= pad_id < vocab_size):
+            raise ValueError(f"pad_id must be in [0, {vocab_size - 1}], got {pad_id}")
+        if not (0 <= dropout_rate <= 1):
+            raise ValueError(f"dropout_rate must be in [0,1], got {dropout_rate}")
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.sequence_length = sequence_length
+        self.pad_id = pad_id
+        self.dropout_rate = dropout_rate
+
+        self.token_embed = nn.Embedding(vocab_size, d_model, padding_idx=pad_id)
+        self.pos_embed = PositionalEmbedding(sequence_length, d_model)
+        self.dropout = nn.Dropout(dropout_rate)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be Tensor, got {type(x)}")
+        if x.dtype != torch.long:
+            raise TypeError(f"x must be torch.long, got {x.dtype}")
+        if x.dim() != 2:
+            raise ValueError(f"x must be 2D (B, S), got shape {tuple(x.shape)}")
+
+        tok = self.token_embed(x)  # (B, S, D) — S can be 0
+        return self.dropout(self.pos_embed(tok))  # (B, S, D)
+
+
+class PositionalEmbedding(nn.Module):
+    """
+    Adds sinusoidal positional encodings.
+
+    Args:
+        sequence_length (int): Maximum sequence length (≥0).
+        d_model (int): Embedding dimension (>0).
+
+    Notes:
+        - Zero-length tables are supported (shape [1, 0, D]).
+    """
+
+    def __init__(self, sequence_length: int, d_model: int):
+        super().__init__()
+
+        if sequence_length < 0:
+            raise ValueError(f"sequence_length must be >= 0, got {sequence_length}")
+        if d_model <= 0:
+            raise ValueError(f"d_model must be > 0, got {d_model}")
+
+        self.sequence_length = sequence_length
+        self.d_model = d_model
+
+        pe = sinusoidal_positional_encoding(sequence_length, d_model)  # [S, D]
+        self.register_buffer("pe", pe.unsqueeze(0))  # [1, S, D]
+
+    def forward(self, x: Tensor) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be Tensor, got {type(x)}")
+        if x.dim() != 3:
+            raise ValueError(f"x must be 3D (B, S, D), got shape {tuple(x.shape)}")
+
+        _, S, D = x.shape
+        if D != self.d_model:
+            raise ValueError(f"d_model mismatch: got {D}, expected {self.d_model}")
+        if S > self.sequence_length:
+            raise ValueError(f"seq_len {S} exceeds max_seq_len {self.sequence_length}")
+
+        return x + self.pe[:, :S].to(dtype=x.dtype, device=x.device)

src/transformer/modules/encoder.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from transformer.modules.attention import MultiHeadAttention
+from transformer.modules.feedforward import FeedForwardLayer
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float):
+        super().__init__()
+
+        # --- type checks (match project convention)
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if not isinstance(num_heads, int):
+            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
+        if not isinstance(d_ff, int):
+            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        # --- value checks (match wording)
+        if not d_model > 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
+        if not num_heads > 0:
+            raise ValueError(f"num_heads must be strictly greater than 0, got {num_heads}")
+        if not d_ff > 0:
+            raise ValueError(f"d_ff must be strictly greater than 0, got {d_ff}")
+        if not (0.0 <= dropout_rate < 1.0):
+            raise ValueError(f"dropout_rate must be between 0 and 1 excluded, got {dropout_rate}")
+
+        self.attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
+        self.feed_forward = FeedForwardLayer(d_model, d_ff)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout2 = nn.Dropout(dropout_rate)
+
+    def forward(self, x: Tensor, src_padding_mask: Tensor | None) -> Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
+        if not (x.dim() == 3):
+            raise ValueError(
+                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+            )
+
+        x = self.norm1(x + self.dropout1(self.attention_layer(x, x, x, src_padding_mask)))
+        x = self.norm2(x + self.dropout2(self.feed_forward(x)))
+        return x
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(
+        self, d_model: int, num_heads: int, d_ff: int, num_layers: int, dropout_rate: float
+    ):
+        super().__init__()
+
+        if not isinstance(num_layers, int):
+            raise TypeError(f"num_layers must be an int, got {type(num_layers)}")
+        if not num_layers > 0:
+            raise ValueError(f"num_layers must be strictly greater than 0, got {num_layers}")
+
+        self.layers = nn.ModuleList(
+            [EncoderLayer(d_model, num_heads, d_ff, dropout_rate) for _ in range(num_layers)]
+        )
+
+    def forward(self, x: Tensor, src_padding_mask: Tensor | None) -> Tensor:
+        for layer in self.layers:
+            x = layer(x, src_padding_mask)
+        return x

src/transformer/modules/feedforward.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+
+
+class FeedForwardLayer(nn.Module):
+    """
+    Position-wise feed-forward layer used in Transformer blocks.
+
+    Architecture:
+        fc1: Linear(d_model -> d_ff)
+        activation: ReLU
+        dropout: nn.Dropout(dropout_rate)
+        fc2: Linear(d_ff -> d_model)
+
+    Args:
+        d_model (int): Dimensionality of model embeddings.
+        d_ff (int): Hidden dimensionality of feed-forward layer.
+        dropout_rate (float): Dropout probability between 0 and 1 (exclusive).
+
+    Shape:
+        Input:  (B, S, D) where D == d_model
+        Output: (B, S, D)
+    """
+
+    def __init__(self, d_model: int, d_ff: int, dropout_rate: float = 0.1):
+        super().__init__()
+
+        if not isinstance(d_ff, int):
+            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if not isinstance(dropout_rate, float):
+            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
+
+        if d_ff <= 0:
+            raise ValueError(f"d_ff must be strictly greater than 0, got {d_ff}")
+        if d_model <= 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
+        if not (0.0 <= dropout_rate < 1.0):
+            raise ValueError(f"dropout_rate must be in [0,1), got {dropout_rate}")
+
+        self.d_model = d_model
+        self.d_ff = d_ff
+        self.dropout_rate = dropout_rate
+
+        self.fc1 = nn.Linear(d_model, d_ff)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(dropout_rate)
+        self.fc2 = nn.Linear(d_ff, d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
+        if x.ndim != 3:
+            raise ValueError(f"x must be 3D of shape (B,S,D); got shape {tuple(x.shape)}")
+        if x.shape[-1] != self.d_model:
+            raise ValueError(f"Last dim {x.shape[-1]} must match d_model {self.d_model}")
+
+        return self.fc2(self.dropout(self.relu(self.fc1(x))))

src/transformer/modules/lm_head.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+
+
+class LMHead(nn.Module):
+    """
+    Language modeling head: projects hidden states to vocabulary logits.
+
+    Args:
+        d_model (int): Model hidden dimension (>0).
+        vocab_size (int): Vocabulary size (>0).
+
+    Input:
+        x (Tensor): shape (B, S, D) with D == d_model. S may be 0.
+
+    Output:
+        logits (Tensor): shape (B, S, V) with V == vocab_size.
+    """
+
+    def __init__(self, d_model: int, vocab_size: int):
+        super().__init__()
+        if not isinstance(vocab_size, int):
+            raise TypeError(f"vocab_size must be an int, got {type(vocab_size)}")
+        if not isinstance(d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(d_model)}")
+        if vocab_size <= 0:
+            raise ValueError(f"vocab_size must be strictly greater than 0, got {vocab_size}")
+        if d_model <= 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
+
+        self.d_model = d_model
+        self.vocab_size = vocab_size
+        self.fc = nn.Linear(d_model, vocab_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
+        if x.dim() != 3:
+            raise ValueError(
+                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+            )
+
+        B, S, D = x.shape
+        if D != self.d_model:
+            raise ValueError(f"Last dim {D} must match d_model {self.d_model}")
+
+        return self.fc(x)

src/transformer/transformer.py

@@ -0,0 +1,239 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from transformer import modules
+from transformer.configs import BasicEncDecCfg
+from transformer.utils import create_causal_mask, sample_from_logits, shift_right
+
+
+class BasicEncoderDecoderTransformer(nn.Module):
+    def __init__(self, cfg: BasicEncDecCfg):
+        super().__init__()
+
+        # Basic type/value guarding for critical cfg fields (mirror project style)
+        if not isinstance(cfg, BasicEncDecCfg):
+            raise TypeError(f"cfg must be a BasicEncDecCfg, got {type(cfg)}")
+        if not isinstance(cfg.vocab_size, int):
+            raise TypeError(f"vocab_size must be an int, got {type(cfg.vocab_size)}")
+        if not cfg.vocab_size > 0:
+            raise ValueError(f"vocab_size must be strictly greater than 0, got {cfg.vocab_size}")
+        if not isinstance(cfg.d_model, int):
+            raise TypeError(f"d_model must be an int, got {type(cfg.d_model)}")
+        if not cfg.d_model > 0:
+            raise ValueError(f"d_model must be strictly greater than 0, got {cfg.d_model}")
+        if not isinstance(cfg.max_seq_len, int):
+            raise TypeError(f"max_seq_len must be an int, got {type(cfg.max_seq_len)}")
+        if not cfg.max_seq_len >= 0:
+            raise ValueError(
+                f"max_seq_len must be greater than or equal to 0, got {cfg.max_seq_len}"
+            )
+        if not isinstance(cfg.pad_id, int):
+            raise TypeError(f"pad_id must be an int, got {type(cfg.pad_id)}")
+        if not (0 <= cfg.pad_id < cfg.vocab_size):
+            raise ValueError(f"pad_id must be in [0, {cfg.vocab_size - 1}], got {cfg.pad_id}")
+        if not isinstance(cfg.bos_id, int):
+            raise TypeError(f"bos_id must be an int, got {type(cfg.bos_id)}")
+        if not (0 <= cfg.bos_id < cfg.vocab_size):
+            raise ValueError(f"bos_id must be in [0, {cfg.vocab_size - 1}], got {cfg.bos_id}")
+        if not isinstance(cfg.eos_id, int):
+            raise TypeError(f"eos_id must be an int, got {type(cfg.eos_id)}")
+        if not (0 <= cfg.eos_id < cfg.vocab_size):
+            raise ValueError(f"eos_id must be in [0, {cfg.vocab_size - 1}], got {cfg.eos_id}")
+
+        self.cfg = cfg
+
+        self.embed = modules.InputEmbedding(
+            cfg.vocab_size, cfg.d_model, cfg.max_seq_len, cfg.pad_id, cfg.dropout_rate
+        )
+        self.encoder = modules.TransformerEncoder(
+            cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.dropout_rate
+        )
+        self.decoder = modules.TransformerDecoder(
+            cfg.d_model, cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.dropout_rate
+        )
+        self.lm_head = modules.LMHead(cfg.d_model, cfg.vocab_size)
+
+        # tie weights
+        self.embed.token_embed.weight = self.lm_head.fc.weight
+
+        # expose a few attrs for convenience
+        self.max_seq_len = cfg.max_seq_len
+        self.pad_id = cfg.pad_id
+        self.bos_id = cfg.bos_id
+        self.eos_id = cfg.eos_id
+
+    def forward(
+        self,
+        src_ids: Tensor,
+        tgt_ids: Tensor,
+        src_padding_mask: Tensor | None,
+        tgt_padding_mask: Tensor | None,
+    ) -> Tensor:
+        # Validate inputs (match project style)
+        if not isinstance(src_ids, torch.Tensor):
+            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
+        if not isinstance(tgt_ids, torch.Tensor):
+            raise TypeError(f"tgt_ids must be a torch.Tensor, got {type(tgt_ids)}")
+        if src_ids.dim() != 2:
+            raise ValueError(
+                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
+            )
+        if tgt_ids.dim() != 2:
+            raise ValueError(
+                f"tgt_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(tgt_ids.shape)}"
+            )
+        if src_ids.dtype != torch.long or tgt_ids.dtype != torch.long:
+            raise TypeError("src_ids and tgt_ids must be torch.long (int64)")
+        if src_padding_mask is not None and not isinstance(src_padding_mask, torch.Tensor):
+            raise TypeError(
+                f"src_padding_mask must be a torch.Tensor or None, got {type(src_padding_mask)}"
+            )
+        if tgt_padding_mask is not None and not isinstance(tgt_padding_mask, torch.Tensor):
+            raise TypeError(
+                f"tgt_padding_mask must be a torch.Tensor or None, got {type(tgt_padding_mask)}"
+            )
+
+        hidden_states = self.encode(src_ids, src_padding_mask)
+        # teacher forcing: shift right with BOS, keep PAD for padding
+        shifted_tgt = shift_right(tgt_ids, self.bos_id, self.pad_id)
+        shifted_tgt_mask = None
+        if tgt_padding_mask is not None:
+            m = tgt_padding_mask.to(dtype=torch.bool, device=shifted_tgt.device)
+
+            # Handle common mask shapes by shifting along the last (time) dimension
+            if m.dim() == 4:  # [B, 1, 1, T]
+                pad_col = torch.zeros_like(m[..., :1])  # [B,1,1,1] False
+                shifted_tgt_mask = torch.cat([pad_col, m[..., :-1]], dim=-1)
+            elif m.dim() == 3:  # [B, 1, T] or [B, H, T]
+                pad_col = torch.zeros_like(m[..., :1])  # [B,*,1] False
+                shifted_tgt_mask = torch.cat([pad_col, m[..., :-1]], dim=-1)
+            elif m.dim() == 2:  # [B, T]
+                pad_col = torch.zeros_like(m[:, :1])  # [B,1] False
+                shifted_tgt_mask = torch.cat([pad_col, m[:, :-1]], dim=-1)
+            else:
+                raise ValueError(f"tgt_padding_mask must have 2, 3, or 4 dims; got {m.dim()}")
+
+        dec_hidden_states = self.decode(
+            hidden_states, shifted_tgt, src_padding_mask, shifted_tgt_mask
+        )
+        logits = self.lm_head(dec_hidden_states)
+        return logits  # (B, T_tgt, V)
+
+    def encode(self, src_ids: Tensor, src_padding_mask: Tensor | None) -> Tensor:
+        if not isinstance(src_ids, torch.Tensor):
+            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
+        if src_ids.dim() != 2:
+            raise ValueError(
+                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
+            )
+        if src_ids.dtype != torch.long:
+            raise TypeError("src_ids must be torch.long (int64)")
+        if src_padding_mask is not None and not isinstance(src_padding_mask, torch.Tensor):
+            raise TypeError(
+                f"src_padding_mask must be a torch.Tensor or None, got {type(src_padding_mask)}"
+            )
+
+        x = self.embed(src_ids)  # (B, Sx, D)
+        hidden_states = self.encoder(x, src_padding_mask)
+        return hidden_states  # (B, Sx, D)
+
+    def decode(
+        self,
+        hidden_states: Tensor,
+        tgt_ids: Tensor,
+        src_padding_mask: Tensor | None,
+        tgt_padding_mask: Tensor | None = None,
+    ) -> Tensor:
+        if not isinstance(hidden_states, torch.Tensor):
+            raise TypeError(f"hidden_states must be a torch.Tensor, got {type(hidden_states)}")
+        if hidden_states.dim() != 3:
+            raise ValueError(
+                f"hidden_states must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(hidden_states.shape)}"
+            )
+        if not isinstance(tgt_ids, torch.Tensor):
+            raise TypeError(f"tgt_ids must be a torch.Tensor, got {type(tgt_ids)}")
+        if tgt_ids.dim() != 2:
+            raise ValueError(
+                f"tgt_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(tgt_ids.shape)}"
+            )
+        if tgt_ids.dtype != torch.long:
+            raise TypeError("tgt_ids must be torch.long (int64)")
+        if src_padding_mask is not None and not isinstance(src_padding_mask, torch.Tensor):
+            raise TypeError(
+                f"src_padding_mask must be a torch.Tensor or None, got {type(src_padding_mask)}"
+            )
+        if tgt_padding_mask is not None and not isinstance(tgt_padding_mask, torch.Tensor):
+            raise TypeError(
+                f"tgt_padding_mask must be a torch.Tensor or None, got {type(tgt_padding_mask)}"
+            )
+
+        # causal mask per sequence length (utils returns broadcastable boolean mask)
+        tgt_causal_mask = create_causal_mask(tgt_ids.size(1)).to(tgt_ids.device)
+        y = self.embed(tgt_ids)  # (B, Sy, D)
+        out = self.decoder(hidden_states, y, src_padding_mask, tgt_padding_mask, tgt_causal_mask)
+        return out  # (B, Sy, D)
+
+    def generate(
+        self,
+        src_ids: Tensor,
+        src_padding_mask: Tensor | None,
+        max_new_tokens: int = 20,
+        temperature: float = 1.0,
+        top_k: int | None = None,
+        top_p: float | None = None,
+        do_sample: bool = False,
+    ) -> Tensor:
+        # validate args
+        if not isinstance(src_ids, torch.Tensor):
+            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
+        if src_ids.dim() != 2:
+            raise ValueError(
+                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
+            )
+        if src_ids.dtype != torch.long:
+            raise TypeError("src_ids must be torch.long (int64)")
+        if src_padding_mask is not None and not isinstance(src_padding_mask, torch.Tensor):
+            raise TypeError(
+                f"src_padding_mask must be a torch.Tensor or None, got {type(src_padding_mask)}"
+            )
+        if not isinstance(max_new_tokens, int):
+            raise TypeError(f"max_new_tokens must be an int, got {type(max_new_tokens)}")
+        if max_new_tokens < 0:
+            raise ValueError(
+                f"max_new_tokens must be greater than or equal to 0, got {max_new_tokens}"
+            )
+        if not isinstance(temperature, float):
+            raise TypeError(f"temperature must be a float, got {type(temperature)}")
+        if not (temperature > 0.0):
+            raise ValueError(f"temperature must be strictly greater than 0, got {temperature}")
+        if top_k is not None and (not isinstance(top_k, int)):
+            raise TypeError(f"top_k must be an int or None, got {type(top_k)}")
+        if top_k is not None and top_k < 0:
+            raise ValueError(f"top_k must be strictly greater than 0, got {top_k}")
+        if top_p is not None and (not isinstance(top_p, float)):
+            raise TypeError(f"top_p must be a float or None, got {type(top_p)}")
+        if top_p is not None and not (0.0 < top_p <= 1.0):
+            raise ValueError(f"top_p must be in (0, 1], got {top_p}")
+        if not isinstance(do_sample, bool):
+            raise TypeError(f"do_sample must be a bool, got {type(do_sample)}")
+        self.eval()
+        with torch.no_grad():
+            batch_size, _ = src_ids.shape
+            device = src_ids.device
+            tgt_ids = torch.full((batch_size, 1), self.bos_id, device=device, dtype=torch.long)
+            finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
+            hidden_states = self.encode(src_ids, src_padding_mask)
+            for _ in range(max_new_tokens):
+                # decode on current tgt_ids; take last-step logits
+                step_hidden = self.decode(hidden_states, tgt_ids, src_padding_mask)  # (B, T, D)
+                logits = self.lm_head(step_hidden)[:, -1, :]  # (B, V)
+                next_ids = sample_from_logits(
+                    logits, temperature=temperature, top_k=top_k, top_p=top_p, do_sample=do_sample
+                )  # (B,)
+                next_ids = torch.where(finished, torch.full_like(next_ids, self.eos_id), next_ids)
+                tgt_ids = torch.cat([tgt_ids, next_ids.unsqueeze(1)], dim=-1)
+                finished |= next_ids == self.eos_id
+                if finished.all():
+                    return tgt_ids
+            return tgt_ids

src/transformer/utils.py

@@ -0,0 +1,612 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable
+
+import torch
+import torch.nn.functional as F
+
+
+def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:
+    """
+    Split the last (model) dimension of a 3D torch.Tensor into (num_heads, d_head) and permute to
+    (batch_size, num_heads, seq_length, d_head).
+
+    Args:
+        x (torch.Tensor): Input torch.Tensor of shape (batch_size, seq_length, d_model).
+        num_heads (int): Number of attention heads to split into. Must be a positive integer
+            that divides d_model exactly.
+
+    Returns:
+        torch.Tensor: torch.Tensor of shape (batch_size, num_heads, seq_length, d_head), where
+            d_head = d_model // num_heads. The output torch.Tensor retains the same dtype and
+            device as the input.
+
+    Invariants:
+        - Output device == x.device
+        - Output dtype == x.dtype
+        - batch_size and seq_length are preserved from the input shape.
+
+    Notes:
+        - Zero-length sequences (seq_length == 0) are supported and will return a torch.Tensor
+          with shape (batch_size, num_heads, 0, d_head).
+
+    Raises:
+        TypeError: If x is not a torch.Tensor or num_heads is not an int.
+        ValueError: If x is not 3D, if num_heads <= 0, or if d_model is not divisible by num_heads.
+    """
+    # Type checks
+    if not isinstance(x, torch.Tensor):
+        raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
+    if not isinstance(num_heads, int):
+        raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
+
+    # Shape checks
+    if x.ndim != 3:
+        raise ValueError(
+            f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
+        )
+    if num_heads <= 0:
+        raise ValueError(f"num_heads must be > 0; got {num_heads}")
+
+    batch_size, seq_length, d_model = x.shape
+
+    # Divisibility check
+    if d_model % num_heads != 0:
+        raise ValueError(
+            f"d_model ({d_model}) must be divisible by num_heads ({num_heads}); "
+            f"got remainder {d_model % num_heads}"
+        )
+
+    d_head = d_model // num_heads
+
+    # Reshape and permute: (B, S, D) -> (B, S, H, Dh) -> (B, H, S, Dh)
+    x = x.reshape(batch_size, seq_length, num_heads, d_head)
+    x = x.permute(0, 2, 1, 3)
+
+    return x
+
+
+def calculate_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    mask: torch.Tensor | None,
+    *,
+    deterministic: bool = False,
+    return_probs: bool = False,
+    attn_dropout_p: float = 0.0,
+):
+    """
+    Scaled dot-product attention with:
+      • device-safe masking (mask auto-moved to query.device),
+      • explicit broadcastability checks for mask,
+      • numerically stable softmax (row-wise max subtraction),
+      • fully-masked rows -> zero probabilities and zero outputs,
+      • fp16/bf16-safe compute via fp32 upcast, output downcast.
+
+    Args:
+        query: (B, H, S_q, D)
+        key:   (B, H, S_kv, D)
+        value: (B, H, S_kv, D)
+        mask: broadcastable to (B, H, S_q, S_kv)
+              - bool: True = masked/ignored position
+              - float: additive bias (e.g., large negative for masked, ALiBi, etc.)
+        deterministic: best-effort reproducibility toggle
+        return_probs: if False, only return attention output
+
+    Returns:
+        attention: (B, H, S_q, D)
+        probs (optional): (B, H, S_q, S_kv)
+    """
+    # ---- Basic shape checks
+    if query.dim() != 4 or key.dim() != 4 or value.dim() != 4:
+        raise ValueError(
+            f"q/k/v must be 4D (B,H,S,D). Got "
+            f"q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}"
+        )
+
+    Bq, Hq, Sq, Dq = query.shape
+    Bk, Hk, Sk, Dk = key.shape
+    Bv, Hv, Sv, Dv = value.shape
+
+    if not (Bq == Bk == Bv):
+        raise ValueError(f"Batch mismatch: q={Bq}, k={Bk}, v={Bv}")
+    if not (Hq == Hk == Hv):
+        raise ValueError(f"Heads mismatch: q={Hq}, k={Hk}, v={Hv}")
+    if Sk != Sv:
+        raise ValueError(f"Key/Value seq length mismatch: Sk={Sk} vs Sv={Sv}")
+    if not (Dq == Dk == Dv):
+        raise ValueError(f"Head dimension mismatch: Dq={Dq}, Dk={Dk}, Dv={Dv}")
+
+    target_shape = (Bq, Hq, Sq, Sk)
+
+    # ---- Device checks: q/k/v must be on the same device (fail fast); mask will be auto-moved
+    if not (query.device == key.device == value.device):
+        raise RuntimeError(
+            f"q/k/v must be on the same device, got "
+            f"query={query.device}, key={key.device}, value={value.device}"
+        )
+
+    # ---- Early exit for empty sequences
+    if Sq == 0:
+        empty_attn = query.new_zeros(Bq, Hq, 0, Dq)
+        empty_probs = query.new_zeros(Bq, Hq, 0, Sk)
+        return (empty_attn, empty_probs) if return_probs else empty_attn
+    if Sk == 0:
+        zero_attn = query.new_zeros(Bq, Hq, Sq, Dq)
+        zero_probs = query.new_zeros(Bq, Hq, Sq, 0)
+        return (zero_attn, zero_probs) if return_probs else zero_attn
+
+    # ---- Compute dtype policy: upcast to fp32 for stability if needed
+    input_dtype = query.dtype
+    compute_dtype = torch.float32 if input_dtype in (torch.float16, torch.bfloat16) else input_dtype
+    q = query.to(compute_dtype)
+    k = key.to(compute_dtype)
+    v = value.to(compute_dtype)
+
+    # ---- Determinism (best effort)
+    if deterministic:
+        try:
+            torch.use_deterministic_algorithms(True)
+        except Exception:
+            pass
+
+    # ---- Scores
+    d_head = Dq
+    scores = torch.matmul(q, k.transpose(-1, -2)) / (d_head**0.5)  # (B,H,Sq,Sk)
+
+    # ---- Mask handling (device + broadcastability + semantics)
+    full_row_mask = None  # (B,H,Sq,1) True where the entire key row is invalid
+    if mask is not None:
+        # Move mask to the same device
+        if mask.device != q.device:
+            mask = mask.to(q.device)
+
+        # Check broadcastability to target shape, then expand (without copy)
+        def _expand_or_error(x: torch.Tensor, name: str) -> torch.Tensor:
+            # Try expanding to target_shape; torch.Tensor.expand raises if not broadcastable
+            try:
+                return x.expand(target_shape)
+            except RuntimeError as e:
+                raise ValueError(
+                    f"{name} with shape {tuple(x.shape)} is not broadcastable to {target_shape}. "
+                    f"Broadcasting rules require dimensions to be equal or 1 in the source."
+                ) from e
+
+        if mask.dtype == torch.bool:
+            mask = _expand_or_error(mask, "Boolean mask")
+            # Use a large finite negative to avoid (-inf) - (-inf) during stabilization
+            neg_large = torch.finfo(compute_dtype).min / 4
+            scores = scores.masked_fill(mask, neg_large)
+            full_row_mask = mask.all(dim=-1, keepdim=True)  # (B,H,Sq,1)
+        else:
+            # Additive bias mask
+            mask = _expand_or_error(mask.to(compute_dtype), "Additive mask")
+            scores = scores + mask
+            # If caller used -inf as additive, detect fully-masked rows after addition
+            full_row_mask = torch.isneginf(scores).all(dim=-1, keepdim=True)
+
+    # ---- Numerically stable softmax: subtract row-wise max
+    row_max = torch.amax(scores, dim=-1, keepdim=True)
+    if full_row_mask is None:
+        # If any rows are -inf (rare with finite negatives), neutralize their max to 0
+        row_max = torch.where(torch.isneginf(row_max), torch.zeros_like(row_max), row_max)
+    else:
+        # For rows marked fully masked, set max to 0 to avoid -inf - (-inf)
+        row_max = torch.where(full_row_mask, torch.zeros_like(row_max), row_max)
+
+    stable_scores = scores - row_max
+    probs = F.softmax(stable_scores, dim=-1)
+
+    # ---- Zero-out fully masked rows (probabilities and thus outputs)
+    if full_row_mask is not None:
+        probs = probs * (~full_row_mask).to(probs.dtype)
+    if attn_dropout_p and attn_dropout_p > 0.0:
+        probs = F.dropout(probs, p=attn_dropout_p, training=True)
+    # ---- Weighted sum with values; cast back to input dtype
+    attention = torch.matmul(probs, v).to(input_dtype)
+
+    if return_probs:
+        out_probs_dtype = (
+            input_dtype if input_dtype in (torch.float16, torch.bfloat16) else probs.dtype
+        )
+        return attention, probs.to(out_probs_dtype)
+    else:
+        return attention
+
+
+def join_heads(x: torch.Tensor) -> torch.Tensor:
+    """
+    Merge multi-head attention outputs into model dimension.
+
+    Args:
+        x (torch.Tensor): shape (batch_size, num_heads, seq_length, d_head)
+
+    Returns:
+        torch.Tensor: shape (batch_size, seq_length, d_model) where d_model = num_heads * d_head
+    """
+    # --- Type checks ---
+    if not isinstance(x, torch.Tensor):
+        raise TypeError(f"Expected torch.Tensor, got {type(x)}")
+
+    if x.ndim != 4:
+        raise ValueError(
+            f"Expected 4D torch.Tensor (batch, num_heads, seq_len, d_head), "
+            f"got shape {tuple(x.shape)} with ndim={x.ndim}"
+        )
+
+    batch_size, num_heads, seq_length, d_head = x.shape
+
+    if not all(
+        isinstance(dim, int) and dim >= 0 for dim in (batch_size, num_heads, seq_length, d_head)
+    ):
+        raise ValueError(f"Invalid shape values: {x.shape}")
+
+    d_model = num_heads * d_head
+
+    # --- Safe reshape ---
+    x = x.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_length, d_model)
+
+    return x
+
+
+def sinusoidal_positional_encoding(
+    seq_len: int,
+    dim: int,
+    *,
+    base: float = 10_000.0,
+    dtype: torch.dtype | None = None,
+    device: torch.device | None = None,
+    offset: int = 0,
+    return_positions: bool = False,
+) -> torch.Tensor:
+    """
+    Build a [seq_len, dim] sinusoidal (trig) positional encoding table:
+      PE[pos, 2i]   = sin(pos / (base^(2i/dim)))
+      PE[pos, 2i+1] = cos(pos / (base^(2i/dim)))
+
+    Design goals:
+    - Numerically stable: uses exp with -log(base) to avoid huge powers.
+    - Dtype/device safe: computes frequencies in float32 and casts at the end;
+      supports fp16/bf16 by upcasting, then downcasting.
+    - Odd dimensions handled: the last column is zero-padded when dim is odd.
+    - Supports an integer `offset` so you can continue encodings for cached KV.
+    - No autograd needed: returns a tensor with requires_grad=False.
+
+    Args:
+      seq_len: length of sequence (>=0).
+      dim: embedding dimension (>=1).
+      base: frequency base (default 10k like Vaswani et al.).
+      dtype: desired dtype of the returned table (defaults to torch.get_default_dtype()).
+      device: desired device of the returned table.
+      offset: position offset to start from (useful for incremental decoding).
+      return_positions: if True, also returns the [seq_len] position indices.
+
+    Returns:
+      pe: [seq_len, dim] tensor on `device` with `dtype`.
+      (pos): optional [seq_len] tensor of positions (if return_positions=True).
+
+    Raises:
+      ValueError if seq_len<0 or dim<=0.
+    """
+    if seq_len < 0:
+        raise ValueError(f"seq_len must be >= 0, got {seq_len}")
+    if dim <= 0:
+        raise ValueError(f"dim must be > 0, got {dim}")
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+
+    # Work in float32 for stability regardless of target dtype.
+    work_dtype = torch.float32
+
+    # Positions [offset .. offset+seq_len-1]
+    # Note: torch.arange is dtype-agnostic; compute in float32 later.
+    pos = torch.arange(offset, offset + seq_len, device=device)
+
+    # Frequencies for even indices: exp(-(log(base) * (2i)/dim))
+    # Avoid base**(2i/dim) directly to reduce overflow/underflow risk.
+    half_dim = dim // 2  # number of sin/cos pairs
+    if half_dim > 0:
+        # [half_dim]
+        exponent = torch.arange(0, half_dim, device=device, dtype=work_dtype)
+        inv_freq = torch.exp(-math.log(base) * (2.0 * exponent) / float(dim))  # [half_dim]
+
+        # Outer product: [seq_len, half_dim]
+        # Compute phase = pos[:, None] * inv_freq[None, :]
+        phase = pos.to(work_dtype).unsqueeze(1) * inv_freq.unsqueeze(0)
+
+        sin_part = torch.sin(phase)
+        cos_part = torch.cos(phase)
+
+        # Interleave sin and cos along the last dim
+        pe_even_odd = torch.stack((sin_part, cos_part), dim=-1).reshape(seq_len, 2 * half_dim)
+        if dim % 2 == 1:
+            # Pad the last column with zeros if dim is odd
+            pad = torch.zeros(seq_len, 1, device=device, dtype=work_dtype)
+            pe_work = torch.cat([pe_even_odd, pad], dim=1)
+        else:
+            pe_work = pe_even_odd
+    else:
+        # dim == 1 case → just a zero column (consistent with odd-dim padding above)
+        pe_work = torch.zeros(seq_len, 1, device=device, dtype=work_dtype)
+
+        # If somehow dim==0 was allowed earlier, we'd have raised.
+
+    # Cast once at the end to the requested dtype.
+    pe = pe_work.to(dtype=dtype)
+
+    # Ensure no gradient tracking (users usually register this as a buffer).
+    pe.requires_grad_(False)
+
+    return (pe, pos.to(device)) if return_positions else pe
+
+
+def combine_masks(mask1, mask2):
+    """
+    Combine two boolean masks with logical OR.
+
+    Args:
+        mask1 (torch.Tensor or None): first mask
+        mask2 (torch.Tensor or None): second mask
+
+    Returns:
+        torch.Tensor or None: combined mask or None if both are None
+
+    Notes:
+        - Both masks must be broadcastable to the same shape.
+        - Returns a boolean tensor if inputs are boolean.
+        - If one mask is None, returns the other unchanged.
+    """
+    if mask1 is not None and mask2 is not None:
+        m1 = mask1.to(dtype=torch.bool)
+        m2 = mask2.to(dtype=torch.bool, device=m1.device)
+        if m1.shape != m2.shape:
+            try:
+                return m1 | m2  # allow broadcasting, but safe
+            except RuntimeError as e:
+                raise ValueError(
+                    f"combine_masks: masks not broadcastable: {m1.shape} vs {m2.shape}"
+                ) from e
+        return m1 | m2
+    elif mask1 is not None:
+        return mask1
+    elif mask2 is not None:
+        return mask2
+    else:
+        return None
+
+
+def shift_right(labels: torch.Tensor, bos_id: int = 1, pad_id: int = 0) -> torch.Tensor:
+    """
+    Build decoder input IDs by shifting target labels to the right and
+    prepending a BOS token. Any labels marked -100 (ignore_index) are
+    converted to PAD *after* shifting so they never appear as inputs.
+
+    Args:
+        labels: Long tensor of shape [B, T] or [T].
+        bos_id: Token id to place at position 0.
+        pad_id: Token id used where labels == -100 after shifting.
+
+    Returns:
+        decoder_input_ids: Long tensor of shape [B, T] (same device as labels).
+
+    """
+    if labels.dim() == 1:
+        labels = labels.unsqueeze(0)  # [1, T]
+    if labels.dim() != 2:
+        raise ValueError(f"shift_right expects 1D or 2D tensor, got shape {tuple(labels.shape)}")
+
+    B, T = labels.shape
+    if T == 0:
+        raise ValueError("shift_right: sequence length must be > 0")
+
+    # Work in long dtype on the same device; do NOT modify `labels` in-place.
+    labels = labels.to(torch.long)
+
+    # Start with PADs everywhere, then copy shifted labels, and set BOS.
+    decoder_input_ids = labels.new_full((B, T), pad_id)
+    decoder_input_ids[:, 1:] = labels[:, :-1]
+    decoder_input_ids[:, 0] = bos_id
+
+    # Any positions that inherited -100 from labels should be PAD instead.
+    # (labels contains -100 for "ignore_index" targets; these should never be inputs)
+    mask_ignore = decoder_input_ids.eq(-100)
+    if mask_ignore.any():
+        decoder_input_ids.masked_fill_(mask_ignore, pad_id)
+
+    return decoder_input_ids
+
+
+def _ensure_2d_logits(logits: torch.Tensor) -> torch.Tensor:
+    """
+    Accepts [B, V] or [..., V]. If rank > 2, flattens leading dims into batch.
+    Returns 2D [B', V] along with a flag and original shape for potential future use.
+    """
+    if logits.dim() == 2:
+        return logits
+    if logits.dim() >= 3:
+        V = logits.size(-1)
+        return logits.reshape(-1, V)
+    raise ValueError(
+        f"sample_from_logits: logits must be at least 2D, got shape {tuple(logits.shape)}"
+    )
+
+
+def _apply_allow_deny_mask(
+    logits: torch.Tensor,
+    *,
+    allowed_tokens: Iterable[int] | None,
+    disallowed_tokens: Iterable[int] | None,
+    filter_value: float,
+) -> torch.Tensor:
+    if allowed_tokens is not None:
+        mask = torch.zeros_like(logits, dtype=torch.bool)
+        idx = torch.tensor(list(allowed_tokens), device=logits.device)
+        idx = idx[(idx >= 0) & (idx < logits.size(-1))]
+        if idx.numel() > 0:
+            mask.index_fill_(-1, idx, True)
+        logits = torch.where(mask, logits, torch.full_like(logits, filter_value))
+    if disallowed_tokens is not None:
+        idx = torch.tensor(list(disallowed_tokens), device=logits.device)
+        idx = idx[(idx >= 0) & (idx < logits.size(-1))]
+        if idx.numel() > 0:
+            logits.index_fill_(-1, idx, filter_value)
+    return logits
+
+
+def _top_k_filtering(
+    logits: torch.Tensor,
+    top_k: int | None,
+    *,
+    min_tokens_to_keep: int,
+    filter_value: float,
+) -> torch.Tensor:
+    if top_k is None or top_k <= 0:
+        return logits
+    k = min(max(top_k, min_tokens_to_keep), logits.size(-1))
+    # threshold per row
+    values, _ = torch.topk(logits, k, dim=-1)
+    thresh = values[..., -1, None]
+    return torch.where(logits < thresh, torch.full_like(logits, filter_value), logits)
+
+
+def _top_p_filtering(
+    logits_scaled: torch.Tensor,
+    probs: torch.Tensor,
+    top_p: float,
+    *,
+    min_tokens_to_keep: int,
+    filter_value: float,
+) -> torch.Tensor:
+    if top_p is None or not (0.0 < top_p < 1.0):
+        return logits_scaled
+    # sort by probability
+    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
+    cum = torch.cumsum(sorted_probs, dim=-1)
+    # tokens to remove: everything after first point where cum > top_p
+    to_remove = cum > top_p
+    # always keep at least min_tokens_to_keep highest-prob tokens
+    to_remove[..., :min_tokens_to_keep] = False
+    # scatter back to original indices
+    scatter_mask = torch.zeros_like(to_remove, dtype=torch.bool).scatter(-1, sorted_idx, to_remove)
+    return torch.where(scatter_mask, torch.full_like(logits_scaled, filter_value), logits_scaled)
+
+
+def sample_from_logits(
+    logits: torch.Tensor,
+    *,
+    do_sample: bool = False,
+    temperature: float = 1.0,
+    top_k: int | None = None,
+    top_p: float | None = None,
+    min_tokens_to_keep: int = 1,
+    allowed_tokens: Iterable[int] | None = None,
+    disallowed_tokens: Iterable[int] | None = None,
+    filter_value: float = -float("inf"),
+    rng: torch.Generator | None = None,
+    return_probs: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """
+    Select token IDs from logits with optional sampling & filtering.
+
+    Inputs:
+      logits: [..., V] or [B, V]; will be treated as 2D [B', V] internally.
+      do_sample: False -> greedy argmax; True -> multinomial sampling.
+      temperature: >0; values <1 sharpen, >1 flatten.
+      top_k: keep only top-k logits (after temperature); None/<=0 disables.
+      top_p: nucleus (0<p<1): keep smallest set with prob mass >= p (after temperature).
+      min_tokens_to_keep: safety floor so we never filter out *everything*.
+      allowed_tokens: iterable of token IDs to allow (others masked out).
+      disallowed_tokens: iterable of token IDs to ban (masked out).
+      filter_value: value used for filtered logits (-inf by default).
+      rng: optional torch.Generator for reproducible sampling.
+      return_probs: if True, also returns the final probabilities used.
+
+    Returns:
+      token_ids: [B'] int64 (and optionally probs [B', V] if return_probs=True).
+    """
+    if temperature <= 0:
+        raise ValueError("temperature must be > 0")
+
+    # Work on a stable float32 copy; keep device and shape
+    logits2d = _ensure_2d_logits(logits).to(dtype=torch.float32)
+
+    # Apply allow/deny lists first (hard mask)
+    logits2d = _apply_allow_deny_mask(
+        logits2d,
+        allowed_tokens=allowed_tokens,
+        disallowed_tokens=disallowed_tokens,
+        filter_value=filter_value,
+    )
+
+    # Scale by temperature before top-k/top-p
+    logits_scaled = logits2d / float(temperature)
+
+    # Top-k first (logit space)
+    logits_scaled = _top_k_filtering(
+        logits_scaled, top_k=top_k, min_tokens_to_keep=min_tokens_to_keep, filter_value=filter_value
+    )
+
+    # Compute probs (needed for top-p and sampling)
+    probs = F.softmax(logits_scaled, dim=-1)
+
+    # Top-p (probability space) — re-mask logits_scaled accordingly
+    if top_p is not None and 0.0 < top_p < 1.0:
+        logits_scaled = _top_p_filtering(
+            logits_scaled,
+            probs,
+            top_p=top_p,
+            min_tokens_to_keep=min_tokens_to_keep,
+            filter_value=filter_value,
+        )
+        probs = F.softmax(logits_scaled, dim=-1)  # recompute after top-p masking
+
+    if do_sample:
+        # multinomial sampling with optional RNG for reproducibility
+        next_ids = torch.multinomial(probs, num_samples=1, replacement=True, generator=rng).squeeze(
+            -1
+        )
+    else:
+        # greedy
+        next_ids = torch.argmax(probs, dim=-1)
+
+    next_ids = next_ids.to(dtype=torch.long, device=logits.device)  # back to original device
+    if return_probs:
+        return next_ids, probs.to(device=logits.device, dtype=probs.dtype)
+    return next_ids
+
+
+def create_causal_mask(
+    max_seq_len: int,
+    *,
+    device: torch.device | None = None,
+    dtype: torch.dtype = torch.bool,
+) -> torch.Tensor:
+    """
+    Build a standard causal mask to prevent attending to future positions.
+
+    Shape: [1, 1, max_seq_len, max_seq_len]
+
+    mask[i, j] = True means position j should be masked when predicting i.
+
+    Args:
+        max_seq_len: length of the target sequence (must be > 0).
+        device: torch device to place the mask on (defaults to cpu).
+        dtype: dtype of mask (default: torch.bool).
+               Usually boolean, but some fused kernels want float with -inf/0.
+
+    Returns:
+        Tensor of shape [1, 1, max_seq_len, max_seq_len].
+    """
+    if max_seq_len <= 0:
+        raise ValueError(f"max_seq_len must be > 0, got {max_seq_len}")
+
+    # upper-triangular (excluding main diagonal)
+    mask = torch.triu(
+        torch.ones((1, 1, max_seq_len, max_seq_len), device=device, dtype=torch.bool),
+        diagonal=1,
+    )
+    if dtype != torch.bool:
+        mask = mask.to(dtype)
+    return mask

tests/units/modules/test_attention.py

@@ -0,0 +1,177 @@
+import pytest
+import torch
+
+from transformer.modules.attention import MultiHeadAttention
+
+
+# -----------------------
+# Helpers
+# -----------------------
+def _rand(B=2, S=5, D=12, device="cpu", dtype=torch.float32):
+    return torch.randn(B, S, D, device=device, dtype=dtype)
+
+
+# =======================
+# Constructor checks
+# =======================
+
+
+def test_ctor_type_checks():
+    with pytest.raises(TypeError):
+        MultiHeadAttention(d_model="64", num_heads=4, dropout_rate=0.1)
+    with pytest.raises(TypeError):
+        MultiHeadAttention(d_model=64, num_heads=4.0, dropout_rate=0.1)
+    with pytest.raises(TypeError):
+        MultiHeadAttention(d_model=64, num_heads=4, dropout_rate="0.1")
+
+
+def test_ctor_value_checks():
+    with pytest.raises(ValueError):
+        MultiHeadAttention(d_model=0, num_heads=4, dropout_rate=0.1)
+    with pytest.raises(ValueError):
+        MultiHeadAttention(d_model=64, num_heads=0, dropout_rate=0.1)
+    with pytest.raises(ValueError):
+        MultiHeadAttention(d_model=63, num_heads=4, dropout_rate=0.1)  # not divisible
+
+
+def test_ctor_happy_path():
+    mha = MultiHeadAttention(64, 8, 0.1)
+    assert mha.d_model == 64 and mha.num_heads == 8 and mha.d_head == 8
+
+
+# =======================
+# Forward: type/shape checks
+# =======================
+
+
+def test_forward_requires_tensors_and_mask_type():
+    mha = MultiHeadAttention(32, 4, 0.1)
+    q, k, v = _rand(2, 5, 32), _rand(2, 5, 32), _rand(2, 5, 32)
+
+    with pytest.raises(TypeError):
+        mha("q", k, v, None)
+    with pytest.raises(TypeError):
+        mha(q, "k", v, None)
+    with pytest.raises(TypeError):
+        mha(q, k, "v", None)
+    with pytest.raises(TypeError):
+        mha(q, k, v, mask="bad")
+
+
+def test_forward_rank_and_lastdim_checks():
+    mha = MultiHeadAttention(32, 4, 0.1)
+    q = torch.randn(2, 5, 32)
+    k = torch.randn(2, 5, 32)
+    v = torch.randn(2, 5, 32)
+
+    with pytest.raises(ValueError):
+        mha(q.unsqueeze(0), k, v, None)  # rank 4 for q
+    with pytest.raises(ValueError):
+        mha(q, k.view(10, 32), v, None)  # rank 2 for k
+    with pytest.raises(ValueError):
+        mha(q[..., :16], k, v, None)  # wrong Dq
+    with pytest.raises(ValueError):
+        mha(q, k[..., :16], v, None)  # wrong Dk
+    with pytest.raises(ValueError):
+        mha(q, k, v[..., :16], None)  # wrong Dv
+
+
+def test_forward_batch_and_seq_mismatch_checks():
+    mha = MultiHeadAttention(24, 3, 0.1)
+    q = torch.randn(2, 5, 24)
+    k = torch.randn(3, 5, 24)  # batch mismatch
+    v = torch.randn(2, 5, 24)  # seq mismatch with k
+
+    with pytest.raises(ValueError):
+        mha(q, k, torch.randn(3, 5, 24), None)  # batch mismatch q vs k/v
+    with pytest.raises(ValueError):
+        mha(q, torch.randn(2, 6, 24), v, None)  # Sk != Sv
+
+
+# =======================
+# Forward: happy path + masks + zero-length
+# =======================
+
+
+@pytest.mark.parametrize("B,S,D,H", [(1, 1, 16, 4), (2, 5, 32, 4), (3, 0, 24, 3)])
+def test_forward_shapes_and_device_dtype(B, S, D, H):
+    device = (
+        torch.device(f"cuda:{torch.cuda.current_device()}")
+        if torch.cuda.is_available()
+        else torch.device("cpu")
+    )
+    mha = MultiHeadAttention(D, H, 0.1).to(device)
+    q = _rand(B, S, D, device=device)
+    k = _rand(B, S, D, device=device)
+    v = _rand(B, S, D, device=device)
+
+    out = mha(q, k, v, mask=None)
+    assert out.shape == (B, S, D)
+    assert out.device == device
+    assert out.dtype == q.dtype
+
+
+def test_boolean_mask_blocks_positions():
+    B, S, D, H = 2, 6, 24, 3
+    mha = MultiHeadAttention(D, H, 0.0)  # no dropout for determinism
+    q = _rand(B, S, D)
+    k = _rand(B, S, D)
+    v = _rand(B, S, D)
+
+    # mask last two keys for all heads/queries
+    mask = torch.zeros(B, 1, 1, S, dtype=torch.bool)
+    mask[..., -2:] = True
+
+    out1 = mha(q, k, v, mask=None)
+    out2 = mha(q, k, v, mask=mask)
+    assert not torch.allclose(out1, out2)
+
+
+def test_additive_mask_supported():
+    B, S, D, H = 2, 5, 32, 4
+    mha = MultiHeadAttention(D, H, 0.0)
+    q, k, v = _rand(B, S, D), _rand(B, S, D), _rand(B, S, D)
+    add = torch.zeros(1, 1, 1, S)
+    out = mha(q, k, v, add)
+    assert out.shape == (B, S, D)
+
+
+# =======================
+# Gradients smoke
+# =======================
+
+
+def test_gradients_flow():
+    mha = MultiHeadAttention(32, 4, 0.1)
+    q = _rand(
+        2,
+        5,
+        32,
+        dtype=torch.float32,
+        device="cpu",
+    )
+    k = _rand(
+        2,
+        5,
+        32,
+        dtype=torch.float32,
+        device="cpu",
+    )
+    v = _rand(
+        2,
+        5,
+        32,
+        dtype=torch.float32,
+        device="cpu",
+    )
+    q.requires_grad_(True)
+    k.requires_grad_(True)
+    v.requires_grad_(True)
+
+    out = mha(q, k, v, mask=None)
+    loss = out.pow(2).mean()
+    loss.backward()
+
+    for t in (q, k, v):
+        assert t.grad is not None
+        assert torch.isfinite(t.grad).all()

tests/units/modules/test_decoder.py

@@ -0,0 +1,94 @@
+import pytest
+import torch
+
+from transformer.modules.decoder import DecoderLayer, TransformerDecoder
+
+
+def _dev():
+    return (
+        torch.device(f"cuda:{torch.cuda.current_device()}")
+        if torch.cuda.is_available()
+        else torch.device("cpu")
+    )
+
+
+# -------- ctor validations --------
+
+
+def test_decoder_layer_param_checks():
+    with pytest.raises(TypeError):
+        DecoderLayer("32", 4, 64, 0.1)
+    with pytest.raises(TypeError):
+        DecoderLayer(32, "4", 64, 0.1)
+    with pytest.raises(TypeError):
+        DecoderLayer(32, 4, "64", 0.1)
+    with pytest.raises(TypeError):
+        DecoderLayer(32, 4, 64, "0.1")
+
+    with pytest.raises(ValueError):
+        DecoderLayer(0, 4, 64, 0.1)
+    with pytest.raises(ValueError):
+        DecoderLayer(32, 0, 64, 0.1)
+    with pytest.raises(ValueError):
+        DecoderLayer(32, 4, 0, 0.1)
+    with pytest.raises(ValueError):
+        DecoderLayer(32, 4, 64, 1.0)  # upper bound excluded
+
+
+def test_transformer_decoder_layers_count_checks():
+    with pytest.raises(TypeError):
+        TransformerDecoder(32, 4, 64, "2", 0.1)
+    with pytest.raises(ValueError):
+        TransformerDecoder(32, 4, 64, 0, 0.1)
+
+
+# -------- forward path --------
+
+
+@pytest.mark.parametrize("B,Sx,Sy,D,H,FF,L", [(2, 5, 6, 24, 3, 48, 2), (1, 0, 0, 16, 1, 32, 1)])
+def test_decoder_forward_happy_path_and_zero_len(B, Sx, Sy, D, H, FF, L):
+    device = _dev()
+    dec = TransformerDecoder(D, H, FF, L, 0.1).to(device)
+
+    x = torch.randn(B, Sx, D, device=device)
+    y = torch.randn(B, Sy, D, device=device)
+
+    src_pad = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device) if Sx > 0 else None
+    tgt_pad = torch.zeros(B, 1, 1, Sy, dtype=torch.bool, device=device) if Sy > 0 else None
+    causal = torch.ones(1, 1, Sy, Sy, dtype=torch.bool, device=device).triu(1) if Sy > 0 else None
+
+    out = dec(x, y, src_pad, tgt_pad, causal)
+    assert out.shape == (B, Sy, D)
+    assert out.device == device
+
+
+def test_decoder_forward_input_checks_and_message_format():
+    layer = DecoderLayer(24, 3, 48, 0.1)
+
+    with pytest.raises(TypeError):
+        layer("not a tensor", torch.randn(2, 3, 24), None, None, None)
+    with pytest.raises(TypeError):
+        layer(torch.randn(2, 3, 24), "not a tensor", None, None, None)
+
+    with pytest.raises(ValueError) as e1:
+        layer(torch.randn(2, 3, 24, 5), torch.randn(2, 3, 24), None, None, None)  # x rank 4
+    assert "x must be a 3D torch.Tensor of shape (B, S, D)" in str(e1.value)
+
+    with pytest.raises(ValueError) as e2:
+        layer(torch.randn(2, 3, 24), torch.randn(2, 3, 24, 5), None, None, None)  # y rank 4
+    assert "y must be a 3D torch.Tensor of shape (B, S, D)" in str(e2.value)
+
+    x = torch.randn(2, 5, 24)
+    y = torch.randn(3, 6, 24)  # batch mismatch
+    with pytest.raises(ValueError) as e3:
+        layer(x, y, None, None, None)
+    assert "Batch size or d_model mismatch between encoder memory and decoder input" in str(
+        e3.value
+    )
+
+    y2 = torch.randn(2, 6, 16)  # d_model mismatch
+    with pytest.raises(ValueError) as e4:
+        layer(x, y2, None, None, None)
+    assert "Batch size or d_model mismatch between encoder memory and decoder input" in str(
+        e4.value
+    )

tests/units/modules/test_embedding.py

@@ -0,0 +1,222 @@
+import math
+
+import pytest
+import torch
+
+from transformer.modules.embedding import InputEmbedding, PositionalEmbedding
+
+
+# -----------------------
+# Helpers
+# -----------------------
+def _rand_ids(B=2, S=5, vocab=11, device="cpu"):
+    return torch.randint(0, vocab, (B, S), dtype=torch.long, device=device)
+
+
+# =======================
+# InputEmbedding — ctor
+# =======================
+
+
+def test_ctor_type_checks():
+    with pytest.raises(TypeError):
+        InputEmbedding("10", 8, 16, 0)  # vocab_size
+    with pytest.raises(TypeError):
+        InputEmbedding(10, "8", 16, 0)  # d_model
+    with pytest.raises(TypeError):
+        InputEmbedding(10, 8, "16", 0)  # sequence_length
+    with pytest.raises(TypeError):
+        InputEmbedding(10, 8, 16, "0")  # pad_id
+
+
+def test_ctor_value_checks_basic():
+    with pytest.raises(ValueError):  # vocab_size <= 0
+        InputEmbedding(0, 8, 16, 0)
+    with pytest.raises(ValueError):  # d_model <= 0
+        InputEmbedding(10, 0, 16, 0)
+    with pytest.raises(ValueError):  # sequence_length < 0 (now allowed to be 0)
+        InputEmbedding(10, 8, -1, 0)
+    with pytest.raises(ValueError):  # pad_id out of range
+        InputEmbedding(10, 8, 16, 10)
+    with pytest.raises(ValueError):
+        InputEmbedding(10, 8, 16, -1)
+
+
+def test_ctor_allows_zero_length_and_sets_attrs():
+    m = InputEmbedding(11, 8, 0, 0)
+    assert m.sequence_length == 0
+
+
+def test_padding_row_zero_and_stays_zero_after_step():
+    m = InputEmbedding(13, 6, 10, pad_id=3)
+    with torch.no_grad():
+        assert torch.allclose(
+            m.token_embed.weight[3], torch.zeros(6, dtype=m.token_embed.weight.dtype)
+        )
+
+    ids = torch.tensor([[3, 4, 5, 6]], dtype=torch.long)  # includes pad token 3
+    out = m(ids).sum()
+    out.backward()
+    opt = torch.optim.SGD(m.parameters(), lr=0.1)
+    opt.step()
+    with torch.no_grad():
+        assert torch.allclose(
+            m.token_embed.weight[3], torch.zeros(6, dtype=m.token_embed.weight.dtype)
+        )
+
+
+# =======================
+# InputEmbedding — forward
+# =======================
+
+
+def test_forward_type_and_shape_checks():
+    m = InputEmbedding(10, 8, 16, 0)
+
+    with pytest.raises(TypeError):
+        m("not a tensor")  # wrong type
+
+    with pytest.raises(TypeError):
+        m(torch.ones(2, 5, dtype=torch.float32))  # wrong dtype, must be long
+
+    with pytest.raises(ValueError):
+        m(torch.ones(2, 5, 1, dtype=torch.long))  # rank != 2
+
+
+def test_forward_out_of_range_ids_raise_index_error():
+    m = InputEmbedding(10, 8, 16, 0)
+    x = torch.tensor([[0, 9, 10]], dtype=torch.long)  # 10 is out of range for vocab_size=10
+    with pytest.raises((IndexError, RuntimeError)):  # PyTorch may raise either
+        m(x)
+
+
+def test_forward_happy_path_shape_dtype_device_and_zero_len():
+    device = (
+        torch.device(f"cuda:{torch.cuda.current_device()}")
+        if torch.cuda.is_available()
+        else torch.device("cpu")
+    )
+    m = InputEmbedding(32, 24, 64, 0).to(device)
+
+    # non-empty
+    x = _rand_ids(B=3, S=7, vocab=32, device=device)
+    out = m(x)
+    assert out.shape == (3, 7, 24)
+    assert out.device == device
+    assert out.dtype == torch.get_default_dtype()
+
+    # zero-length sequence allowed
+    x0 = _rand_ids(B=2, S=0, vocab=32, device=device)
+    out0 = m(x0)
+    assert out0.shape == (2, 0, 24)
+
+
+def test_forward_adds_positions_not_just_tokens():
+    m = InputEmbedding(20, 12, 32, 0)
+    x = _rand_ids(B=2, S=5, vocab=20)
+    tok_only = m.token_embed(x)
+    out = m(x)
+    assert not torch.allclose(out, tok_only)
+
+
+def test_forward_respects_sequence_length_limit():
+    m = InputEmbedding(16, 8, 5, 0)
+    x_ok = _rand_ids(B=2, S=5, vocab=16)
+    _ = m(x_ok)  # should not raise
+    x_bad = _rand_ids(B=2, S=6, vocab=16)
+    with pytest.raises(ValueError) as ei:
+        _ = m(x_bad)
+    assert "exceeds max_seq_len" in str(ei.value)
+
+
+# =======================
+# PositionalEmbedding — ctor
+# =======================
+
+
+def test_positional_ctor_value_checks():
+    with pytest.raises(ValueError):
+        PositionalEmbedding(-1, 8)  # sequence_length < 0
+    with pytest.raises(ValueError):
+        PositionalEmbedding(16, 0)  # d_model <= 0
+
+
+def test_positional_buffer_registered_and_constant_zero_ok():
+    pe = PositionalEmbedding(0, 10)  # zero supported
+    assert hasattr(pe, "pe")
+    assert isinstance(pe.pe, torch.Tensor)
+    assert pe.pe.shape == (1, 0, 10)
+    assert pe.pe.requires_grad is False
+
+
+# =======================
+# PositionalEmbedding — forward
+# =======================
+
+
+def test_positional_forward_type_and_shape_checks():
+    pe = PositionalEmbedding(16, 8)
+
+    with pytest.raises(TypeError):
+        pe("not a tensor")
+    with pytest.raises(ValueError):
+        pe(torch.zeros(2, 5))  # rank != 3
+    with pytest.raises(ValueError):
+        pe(torch.zeros(2, 4, 6))  # d_model mismatch
+    with pytest.raises(ValueError):
+        pe(torch.zeros(2, 17, 8))  # seq_len exceeds max
+
+
+def test_positional_forward_adds_nonzero_positions_and_preserves_shape():
+    pe = PositionalEmbedding(32, 12)
+    x = torch.zeros(2, 7, 12)
+    y = pe(x)
+    assert y.shape == x.shape
+    assert not torch.allclose(y, x)
+
+
+def test_positional_forward_device_and_dtype_follow_input_and_zero_len():
+    pe = PositionalEmbedding(64, 16)
+
+    # CPU float32
+    x = torch.zeros(1, 5, 16, dtype=torch.float32, device="cpu")
+    y = pe(x)
+    assert y.device.type == "cpu" and y.dtype == torch.float32
+
+    # zero length
+    x0 = torch.zeros(1, 0, 16)
+    y0 = pe(x0)
+    assert y0.shape == (1, 0, 16)
+
+    # bfloat16 on CPU (if supported)
+    try:
+        x_bf = torch.zeros(1, 5, 16, dtype=torch.bfloat16)
+        y_bf = pe(x_bf)
+        assert y_bf.dtype == torch.bfloat16
+    except Exception:
+        pass
+
+    # CUDA & half if available
+    if torch.cuda.is_available():
+        xh = torch.zeros(1, 5, 16, dtype=torch.float16, device="cuda")
+        yh = pe(xh)
+        assert yh.device.type == "cuda" and yh.dtype == torch.float16
+
+
+def test_positional_matches_known_small_reference():
+    # Cross-check first few values with textbook sin/cos
+    S, D = 4, 6
+    pe = PositionalEmbedding(S, D)
+    x = torch.zeros(1, S, D)
+    y = pe(x)
+    added = y[0]  # [S, D]
+
+    ref = torch.zeros_like(added)
+    base = 10_000.0
+    for pos in range(S):
+        for i in range(D // 2):
+            denom = base ** (2 * i / D)
+            ref[pos, 2 * i] = math.sin(pos / denom)
+            ref[pos, 2 * i + 1] = math.cos(pos / denom)
+
+    assert torch.allclose(added, ref, atol=1e-6, rtol=1e-6)

tests/units/modules/test_encoder.py

@@ -0,0 +1,69 @@
+import pytest
+import torch
+
+from transformer.modules.encoder import EncoderLayer, TransformerEncoder
+
+
+def _dev():
+    return (
+        torch.device(f"cuda:{torch.cuda.current_device()}")
+        if torch.cuda.is_available()
+        else torch.device("cpu")
+    )
+
+
+# -------- ctor validations --------
+
+
+def test_encoder_layer_param_checks():
+    with pytest.raises(TypeError):
+        EncoderLayer("32", 4, 64, 0.1)
+    with pytest.raises(TypeError):
+        EncoderLayer(32, "4", 64, 0.1)
+    with pytest.raises(TypeError):
+        EncoderLayer(32, 4, "64", 0.1)
+    with pytest.raises(TypeError):
+        EncoderLayer(32, 4, 64, "0.1")
+
+    with pytest.raises(ValueError):
+        EncoderLayer(0, 4, 64, 0.1)
+    with pytest.raises(ValueError):
+        EncoderLayer(32, 0, 64, 0.1)
+    with pytest.raises(ValueError):
+        EncoderLayer(32, 4, 0, 0.1)
+    with pytest.raises(ValueError):
+        EncoderLayer(32, 4, 64, 1.0)  # upper bound excluded
+
+
+def test_transformer_encoder_layers_count_checks():
+    with pytest.raises(TypeError):
+        TransformerEncoder(32, 4, 64, "2", 0.1)
+    with pytest.raises(ValueError):
+        TransformerEncoder(32, 4, 64, 0, 0.1)
+
+
+# -------- forward path --------
+
+
+@pytest.mark.parametrize("B,S,D,H,FF,L", [(2, 5, 24, 3, 48, 2), (1, 0, 16, 1, 32, 1)])
+def test_encoder_forward_happy_path_and_zero_len(B, S, D, H, FF, L):
+    device = _dev()
+    enc = TransformerEncoder(D, H, FF, L, 0.1).to(device)
+
+    x = torch.randn(B, S, D, device=device)
+    src_pad = torch.zeros(B, 1, 1, S, dtype=torch.bool, device=device) if S > 0 else None
+
+    out = enc(x, src_pad)
+    assert out.shape == (B, S, D)
+    assert out.device == device
+
+
+def test_encoder_forward_input_checks_and_message_format():
+    layer = EncoderLayer(24, 3, 48, 0.1)
+
+    with pytest.raises(TypeError):
+        layer("not a tensor", None)
+
+    with pytest.raises(ValueError) as ei:
+        layer(torch.randn(2, 3, 4, 5), None)  # rank 4
+    assert "x must be a 3D torch.Tensor of shape (B, S, D)" in str(ei.value)

tests/units/modules/test_feedforward.py

@@ -0,0 +1,131 @@
+import pytest
+import torch
+
+from transformer.modules.feedforward import FeedForwardLayer
+
+# =======================
+# Constructor checks
+# =======================
+
+
+def test_ctor_type_checks():
+    with pytest.raises(TypeError):
+        FeedForwardLayer("64", 256, 0.1)  # d_model
+    with pytest.raises(TypeError):
+        FeedForwardLayer(64, "256", 0.1)  # d_ff
+    with pytest.raises(TypeError):
+        FeedForwardLayer(64, 256, "0.1")  # dropout_rate
+
+
+def test_ctor_value_checks():
+    with pytest.raises(ValueError):
+        FeedForwardLayer(0, 256, 0.1)
+    with pytest.raises(ValueError):
+        FeedForwardLayer(64, 0, 0.1)
+    with pytest.raises(ValueError):
+        FeedForwardLayer(64, 256, -1.0)
+    with pytest.raises(ValueError):
+        FeedForwardLayer(64, 256, 1.0)  # upper bound excluded
+
+
+def test_ctor_happy_path_defaults_and_attrs():
+    ffn = FeedForwardLayer(64, 256)  # default dropout=0.1
+    assert ffn.d_model == 64 and ffn.d_ff == 256
+    assert isinstance(ffn.dropout, torch.nn.Dropout)
+    assert abs(ffn.dropout.p - 0.1) < 1e-9
+
+
+# =======================
+# Forward: input validation
+# =======================
+
+
+def test_forward_type_and_rank_checks():
+    ffn = FeedForwardLayer(32, 64, 0.1)
+    with pytest.raises(TypeError):
+        ffn("not a tensor")
+    with pytest.raises(ValueError):
+        ffn(torch.randn(5, 32))  # rank 2
+    with pytest.raises(ValueError):
+        ffn(torch.randn(2, 3, 16))  # D mismatch (expects 32)
+
+
+# =======================
+# Forward: shapes, device, dtype, zero-length
+# =======================
+
+
+@pytest.mark.parametrize("B,S,D", [(1, 1, 16), (2, 7, 32), (3, 0, 24)])
+def test_forward_shapes_device_dtype_and_zero_len(B, S, D):
+    device = (
+        torch.device(f"cuda:{torch.cuda.current_device()}")
+        if torch.cuda.is_available()
+        else torch.device("cpu")
+    )
+    ffn = FeedForwardLayer(D, D * 2, 0.2).to(device)
+
+    x = torch.randn(B, S, D, device=device, dtype=torch.float32)
+    y = ffn(x)
+    assert y.shape == (B, S, D)
+    assert y.device == device
+    assert y.dtype == x.dtype
+
+
+# =======================
+# Forward: gradient flow
+# =======================
+
+
+def test_gradients_flow():
+    D, H = 32, 64
+    ffn = FeedForwardLayer(D, H, 0.1)
+    x = torch.randn(2, 5, D, requires_grad=True)
+    y = ffn(x)
+    loss = y.pow(2).mean()
+    loss.backward()
+    assert x.grad is not None and torch.isfinite(x.grad).all()
+
+
+# =======================
+# Dropout behavior
+# =======================
+
+
+def test_dropout_zero_equals_no_dropout_path():
+    D, H = 16, 32
+    # Build two modules: one with p=0.0, one with p>0 but switched to eval()
+    ffn0 = FeedForwardLayer(D, H, 0.0)
+    ffnx = FeedForwardLayer(D, H, 0.5).eval()  # eval disables dropout
+
+    ffnx.load_state_dict(ffn0.state_dict(), strict=False)
+
+    x = torch.randn(2, 4, D)
+    y0 = ffn0(x)
+    yx = ffnx(x)
+    # With dropout disabled both should be equal (same weights distrib not identical, but same computation tree)
+    assert torch.allclose(y0, yx, atol=1e-6, rtol=1e-6)
+
+
+def test_dropout_training_changes_output_vs_eval():
+    D, H = 32, 64
+    ffn = FeedForwardLayer(D, H, 0.5)
+    x = torch.randn(3, 6, D)
+
+    torch.manual_seed(123)
+    ffn.train()
+    y_train = ffn(x)
+
+    torch.manual_seed(123)
+    ffn.eval()
+    y_eval = ffn(x)
+
+    # Same seed but eval disables dropout -> outputs should differ
+    assert not torch.allclose(y_train, y_eval)
+
+
+def test_dropout_is_noop_on_zero_length():
+    D, H = 32, 64
+    ffn = FeedForwardLayer(D, H, 0.7).train()
+    x = torch.randn(2, 0, D)  # zero-length sequence
+    y = ffn(x)
+    assert y.shape == (2, 0, D)  # no crash; shape preserved

tests/units/modules/test_lm_head.py

@@ -0,0 +1,73 @@
+import pytest
+import torch
+
+from transformer.modules.lm_head import LMHead
+
+# -----------------------
+# Constructor checks
+# -----------------------
+
+
+def test_ctor_type_and_value_checks():
+    with pytest.raises(TypeError):
+        LMHead("64", 10)
+    with pytest.raises(TypeError):
+        LMHead(64, "10")
+    with pytest.raises(ValueError):
+        LMHead(0, 10)
+    with pytest.raises(ValueError):
+        LMHead(64, 0)
+
+
+def test_ctor_happy_path():
+    lm = LMHead(32, 1000)
+    assert lm.d_model == 32
+    assert lm.vocab_size == 1000
+    assert lm.fc.weight.shape == (1000, 32)
+
+
+# -----------------------
+# Forward checks (3D only)
+# -----------------------
+
+
+def test_forward_type_and_shape_errors():
+    lm = LMHead(16, 50)
+    with pytest.raises(TypeError):
+        lm("not a tensor")
+    with pytest.raises(ValueError):
+        lm(torch.randn(2, 3, 4, 5))  # rank 4 not supported
+    with pytest.raises(ValueError):
+        lm(torch.randn(5, 16))  # rank 2 not supported
+    with pytest.raises(ValueError):
+        lm(torch.randn(2, 8, 15))  # last dim != d_model
+
+
+def test_forward_happy_path_and_zero_len():
+    device = (
+        torch.device(f"cuda:{torch.cuda.current_device()}")
+        if torch.cuda.is_available()
+        else torch.device("cpu")
+    )
+    lm = LMHead(24, 101).to(device)
+
+    # Non-empty
+    x = torch.randn(3, 7, 24, device=device, dtype=torch.float32)
+    y = lm(x)
+    assert y.shape == (3, 7, 101)
+    assert y.device == device and y.dtype == torch.float32
+
+    # Zero-length sequence allowed
+    x0 = torch.randn(2, 0, 24, device=device, dtype=torch.float32)
+    y0 = lm(x0)
+    assert y0.shape == (2, 0, 101)
+
+
+def test_gradients_flow():
+    lm = LMHead(8, 40)
+    x = torch.randn(3, 6, 8, requires_grad=True)
+    y = lm(x)
+    loss = y.pow(2).mean()
+    loss.backward()
+    assert x.grad is not None
+    assert torch.isfinite(x.grad).all()

tests/units/test_transformer.py

@@ -0,0 +1,326 @@
+import pytest
+import torch
+
+from transformer.configs import BasicEncDecCfg
+from transformer.transformer import BasicEncoderDecoderTransformer
+
+
+def _dev():
+    return (
+        torch.device(f"cuda:{torch.cuda.current_device()}")
+        if torch.cuda.is_available()
+        else torch.device("cpu")
+    )
+
+
+def _tiny_cfg(**over):
+    # Small, fast config for tests; zero-length safe via max_seq_len >= 0
+    base = dict(
+        vocab_size=32,
+        d_model=16,
+        d_ff=32,
+        num_heads=4,
+        num_layers=2,
+        max_seq_len=32,
+        pad_id=0,
+        bos_id=1,
+        eos_id=2,
+        dropout_rate=0.1,
+    )
+    base.update(over)
+    return BasicEncDecCfg(**base)
+
+
+# -----------------------
+# Constructor & config validation
+# -----------------------
+
+
+def test_ctor_cfg_type_and_values():
+    with pytest.raises(TypeError):
+        BasicEncoderDecoderTransformer(cfg="not a cfg")
+    # invalid vocab / dims / indices
+    bads = [
+        dict(vocab_size=0),
+        dict(d_model=0),
+        dict(max_seq_len=-1),
+        dict(pad_id=-1),
+        dict(pad_id=100),  # out of range
+        dict(bos_id=999),
+        dict(eos_id=999),
+    ]
+    for upd in bads:
+        with pytest.raises((TypeError, ValueError)):
+            BasicEncoderDecoderTransformer(_tiny_cfg(**upd))
+
+    # happy path
+    model = BasicEncoderDecoderTransformer(_tiny_cfg())
+    # weight tying sanity: same storage
+    assert model.embed.token_embed.weight.data_ptr() == model.lm_head.fc.weight.data_ptr()
+
+
+# -----------------------
+# Encode/Decode/Forward — shapes, dtypes, devices, zero-length
+# -----------------------
+
+
+@pytest.mark.parametrize("Sx,Sy", [(5, 6), (0, 1)])
+def test_forward_pipeline_shapes_and_types(Sx, Sy):
+    device = _dev()
+    cfg = _tiny_cfg(max_seq_len=16)
+    model = BasicEncoderDecoderTransformer(cfg).to(device)
+
+    B = 2
+    src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
+    tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
+
+    # padding masks (boolean), broadcastable; allow None when S == 0
+    src_mask = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device) if Sx > 0 else None
+    tgt_mask = torch.zeros(B, 1, 1, Sy, dtype=torch.bool, device=device) if Sy > 0 else None
+
+    logits = model(src, tgt, src_mask, tgt_mask)
+    assert logits.shape == (B, Sy, cfg.vocab_size)
+    assert logits.device == device
+    assert logits.dtype == torch.get_default_dtype()
+
+    # encode / decode individually
+    mem = model.encode(src, src_mask)
+    assert mem.shape == (B, Sx, cfg.d_model)
+    out = model.decode(mem, tgt, src_mask, tgt_mask)
+    assert out.shape == (B, Sy, cfg.d_model)
+
+
+def test_forward_dtype_checks_and_shape_errors():
+    cfg = _tiny_cfg()
+    model = BasicEncoderDecoderTransformer(cfg)
+
+    B, Sx, Sy = 2, 4, 5
+    src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long)
+    tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long)
+
+    with pytest.raises(TypeError):
+        model("not a tensor", tgt, None, None)
+    with pytest.raises(TypeError):
+        model(src, "not a tensor", None, None)
+    with pytest.raises(ValueError):
+        model(src.unsqueeze(0), tgt, None, None)  # rank 3 src_ids
+    with pytest.raises(ValueError):
+        model(src, tgt.unsqueeze(-1), None, None)  # rank 3 tgt_ids
+    with pytest.raises(TypeError):
+        model(src.float(), tgt, None, None)  # wrong dtype
+    with pytest.raises(TypeError):
+        model(src, tgt.int(), None, None)  # wrong dtype
+    with pytest.raises(TypeError):
+        model(src, tgt, src_padding_mask="not a tensor", tgt_padding_mask=None)
+    with pytest.raises(TypeError):
+        model(src, tgt, src_padding_mask=None, tgt_padding_mask="not a tensor")
+
+
+def test_encode_decode_errors_and_messages():
+    cfg = _tiny_cfg()
+    model = BasicEncoderDecoderTransformer(cfg)
+
+    src = torch.randint(0, cfg.vocab_size, (2, 4))
+    with pytest.raises(ValueError):
+        model.encode(src.unsqueeze(-1), None)  # rank 3 not allowed
+    with pytest.raises(TypeError):
+        model.encode(src.float(), None)  # dtype must be long
+
+    mem = torch.randn(2, 4, cfg.d_model)
+    tgt = torch.randint(0, cfg.vocab_size, (2, 5))
+    with pytest.raises(ValueError):
+        model.decode(mem.unsqueeze(0), tgt, None, None)  # mem rank 4
+    with pytest.raises(ValueError):
+        model.decode(mem, tgt.unsqueeze(-1), None, None)  # tgt rank 3
+    with pytest.raises(TypeError):
+        model.decode(mem, tgt.float(), None, None)  # tgt dtype must be long
+
+
+def test_exceeding_max_seq_len_raises():
+    cfg = _tiny_cfg(max_seq_len=4)
+    model = BasicEncoderDecoderTransformer(cfg)
+    # src longer than max -> PositionalEmbedding should raise
+    src = torch.randint(0, cfg.vocab_size, (2, 5), dtype=torch.long)
+    with pytest.raises(ValueError):
+        model.encode(src, None)
+
+
+# -----------------------
+# Gradients & weight tying
+# -----------------------
+
+
+def test_backward_through_full_forward_and_tied_weights():
+    device = _dev()
+    cfg = _tiny_cfg()
+    model = BasicEncoderDecoderTransformer(cfg).to(device)
+
+    B, Sx, Sy = 2, 6, 5
+    src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
+    tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
+    src_mask = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device)
+    tgt_mask = torch.zeros(B, 1, 1, Sy, dtype=torch.bool, device=device)
+
+    logits = model(src, tgt, src_mask, tgt_mask)  # (B, Sy, V)
+    # Dummy labels (language modeling): predict tgt itself; ignore pads via mask
+    loss = logits.pow(2).mean()
+    loss.backward()
+
+    # Tied weights should have gradient (one shared tensor)
+    tied_weight = model.lm_head.fc.weight
+    assert tied_weight.grad is not None
+    assert torch.isfinite(tied_weight.grad).all()
+
+
+# -----------------------
+# Generate API
+# -----------------------
+
+
+def test_generate_arg_checks_and_output_shapes():
+    device = _dev()
+    cfg = _tiny_cfg()
+    model = BasicEncoderDecoderTransformer(cfg).to(device)
+
+    B, Sx = 2, 4
+    src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
+    src_mask = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device)
+
+    # type/value errors
+    with pytest.raises(TypeError):
+        model.generate("not", src_mask)
+    with pytest.raises(ValueError):
+        model.generate(src.unsqueeze(-1), src_mask)
+    with pytest.raises(TypeError):
+        model.generate(src.float(), src_mask)
+    with pytest.raises(TypeError):
+        model.generate(src, "not a tensor")
+    with pytest.raises(TypeError):
+        model.generate(src, src_mask, max_new_tokens="10")
+    with pytest.raises(ValueError):
+        model.generate(src, src_mask, max_new_tokens=-1)
+    with pytest.raises(TypeError):
+        model.generate(src, src_mask, temperature="1.0")
+    with pytest.raises(ValueError):
+        model.generate(src, src_mask, temperature=0.0)
+    with pytest.raises(TypeError):
+        model.generate(src, src_mask, top_k="5")
+    with pytest.raises(ValueError):
+        model.generate(src, src_mask, top_k=0)
+    with pytest.raises(TypeError):
+        model.generate(src, src_mask, top_p="0.9")
+    with pytest.raises(ValueError):
+        model.generate(src, src_mask, top_p=0.0)
+    with pytest.raises(ValueError):
+        model.generate(src, src_mask, top_p=1.1)
+
+    # happy path: shapes & dtypes; content is stochastic so we don't assert exact ids
+    out = model.generate(src, src_mask, max_new_tokens=7, temperature=1.0, top_k=None, top_p=None)
+    # Should start with BOS and append new tokens; total length = 1 + new tokens (or earlier if EOS)
+    assert out.shape[0] == B and out.dim() == 2
+    assert out.dtype == torch.long and out.device == device
+    assert out.size(1) >= 1 and out.size(1) <= 1 + 7  # early stop on EOS allowed
+
+    # Ensure BOS is at position 0 for every sequence
+    assert torch.all(out[:, 0] == cfg.bos_id)
+
+
+def test_generate_zero_new_tokens_returns_only_bos():
+    device = _dev()
+    cfg = _tiny_cfg()
+    model = BasicEncoderDecoderTransformer(cfg).to(device)
+
+    src = torch.randint(0, cfg.vocab_size, (2, 3), dtype=torch.long, device=device)
+    out = model.generate(src, None, max_new_tokens=0)
+    assert out.shape == (2, 1)  # just the BOS token
+
+
+# -----------------------
+# Mask broadcast sanity (boolean & additive are handled in utils; we just pass through)
+# -----------------------
+
+
+def test_forward_accepts_boolean_padding_masks():
+    device = _dev()
+    cfg = _tiny_cfg()
+    model = BasicEncoderDecoderTransformer(cfg).to(device)
+
+    B, Sx, Sy = 2, 5, 6
+    src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
+    tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
+
+    src_pad_bool = torch.zeros(B, 1, 1, Sx, dtype=torch.bool, device=device)
+    tgt_pad_bool = torch.zeros(B, 1, 1, Sy, dtype=torch.bool, device=device)
+
+    logits = model(src, tgt, src_pad_bool, tgt_pad_bool)
+    assert logits.shape == (B, Sy, cfg.vocab_size)
+
+
+# -----------------------
+# Weight tying and optimizer step
+# -----------------------
+
+
+def test_weight_tying_and_optimizer_step_changes_weights_once():
+    device = _dev()
+    cfg = _tiny_cfg()
+    model = BasicEncoderDecoderTransformer(cfg).to(device)
+
+    # Same storage pointer means truly tied
+    assert model.embed.token_embed.weight.data_ptr() == model.lm_head.fc.weight.data_ptr()
+
+    opt = torch.optim.SGD(model.parameters(), lr=0.01)
+    B, Sx, Sy = 2, 5, 4
+    src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
+    tgt = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
+    logits = model(src, tgt, None, None)  # (B, Sy, V)
+
+    # Simple loss to get gradients flowing
+    loss = logits.pow(2).mean()
+    loss.backward()
+
+    # The tied tensor should have a gradient and be finite
+    tied = model.lm_head.fc.weight
+    assert tied.grad is not None and torch.isfinite(tied.grad).all()
+
+    # Save copy before step and ensure it changes after step (once, not twice)
+    before = tied.detach().clone()
+    opt.step()
+    after = tied.detach().clone()
+
+    assert not torch.allclose(before, after)
+    # Still tied after step (same storage)
+    assert model.embed.token_embed.weight.data_ptr() == model.lm_head.fc.weight.data_ptr()
+
+
+# -----------------------
+# Autoregressive invariants: causal mask honored
+# -----------------------
+
+
+def test_decode_invariance_to_future_tokens():
+    """
+    Sanity-check causal masking: logits at position t must not depend on tokens at > t.
+    Construct two target sequences that are identical up to t and differ after; compare
+    decoder outputs at step t.
+    """
+    device = _dev()
+    cfg = _tiny_cfg()
+    model = BasicEncoderDecoderTransformer(cfg).to(device).eval()
+
+    B, Sx, Sy = 1, 4, 5
+    src = torch.randint(0, cfg.vocab_size, (B, Sx), dtype=torch.long, device=device)
+    mem = model.encode(src, None)
+
+    # Construct y1 and y2 identical up to t=3, different afterwards
+    y1 = torch.randint(0, cfg.vocab_size, (B, Sy), dtype=torch.long, device=device)
+    y2 = y1.clone()
+    t = 3
+    if Sy > t + 1:
+        y2[:, t + 1 :] = (y1[:, t + 1 :] + 1) % cfg.vocab_size
+
+    out1 = model.decode(mem, y1, None, None)  # (B, Sy, D)
+    out2 = model.decode(mem, y2, None, None)
+
+    # Compare hidden states up to and including position t
+    assert torch.allclose(out1[:, : t + 1], out2[:, : t + 1], atol=1e-5, rtol=1e-5)

tests/units/test_utils.py

@@ -0,0 +1,749 @@
+import math
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from transformer.utils import (
+    calculate_attention,
+    combine_masks,
+    create_causal_mask,
+    join_heads,
+    sample_from_logits,
+    shift_right,
+    sinusoidal_positional_encoding,
+    split_heads,
+)
+
+###___split_heads___###
+
+
+def test_happy_path_shape_dtype_device_cpu():
+    x = torch.randn(2, 10, 16)  # B=2, S=10, D=16
+    out = split_heads(x, num_heads=4)
+    assert out.shape == (2, 4, 10, 4)
+    assert out.dtype == x.dtype
+    assert out.device == x.device
+
+
+def test_smallest_valid():
+    x = torch.randn(1, 1, 4)
+    out = split_heads(x, num_heads=2)
+    assert out.shape == (1, 2, 1, 2)
+
+
+def test_zero_length_seq_allowed():
+    x = torch.randn(2, 0, 8)
+    out = split_heads(x, num_heads=2)
+    assert out.shape == (2, 2, 0, 4)
+
+
+def test_non_3d_input_raises_value_error():
+    x2 = torch.randn(10, 16)
+    with pytest.raises(ValueError) as ei2:
+        split_heads(x2, 4)
+    assert "(B, S, D)" in str(ei2.value)
+
+    x4 = torch.randn(2, 3, 4, 5)
+    with pytest.raises(ValueError) as ei4:
+        split_heads(x4, 4)
+    assert "(B, S, D)" in str(ei4.value)
+
+
+def test_invalid_num_heads_type_raises_type_error():
+    x = torch.randn(2, 10, 16)
+    with pytest.raises(TypeError):
+        split_heads(x, 4.0)  # float
+    with pytest.raises(TypeError):
+        split_heads(x, torch.tensor(4))  # Tensor
+
+
+def test_num_heads_leq_zero_raises_value_error():
+    x = torch.randn(2, 10, 16)
+    with pytest.raises(ValueError):
+        split_heads(x, 0)
+    with pytest.raises(ValueError):
+        split_heads(x, -2)
+
+
+def test_non_divisible_d_model_raises_value_error_message_has_values():
+    x = torch.randn(2, 5, 10)
+    with pytest.raises(ValueError) as e:
+        split_heads(x, 3)
+    msg = str(e.value)
+    assert "d_model (10)" in msg and "num_heads (3)" in msg
+
+
+@pytest.mark.parametrize(
+    "B,S,D,H",
+    [
+        (1, 1, 8, 1),
+        (2, 3, 12, 3),
+        (4, 7, 32, 8),
+        (3, 0, 24, 6),
+    ],
+)
+def test_invariants_random_valid(B, S, D, H):
+    x = torch.randn(B, S, D)
+    out = split_heads(x, H)
+    assert out.shape[0] == B
+    assert out.shape[2] == S
+    assert (H * (D // H)) == D
+    assert x.numel() == out.numel()
+
+
+def test_grad_propagates():
+    x = torch.randn(2, 10, 16, requires_grad=True)
+    out = split_heads(x, 4)
+    # Do a simple scalar reduction to make grad non-trivial
+    loss = out.square().mean()
+    loss.backward()
+    assert x.grad is not None
+    assert x.grad.shape == x.shape
+
+
+def test_non_contiguous_input():
+    # Create a non-contiguous view by transposing and slicing
+    base = torch.randn(10, 2, 16)
+    x = base.transpose(0, 1)[:, :8, :]  # shape (2, 8, 16), likely non-contiguous
+    assert not x.is_contiguous()
+    out = split_heads(x, 4)
+    assert out.shape == (2, 4, 8, 4)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_device_preserved_cuda():
+    x = torch.randn(2, 10, 16, device="cuda")
+    out = split_heads(x, 4)
+    assert out.device.type == "cuda"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_dtype_half_bfloat_on_gpu():
+    x_half = torch.randn(2, 10, 16, device="cuda", dtype=torch.float16)
+    out_half = split_heads(x_half, 4)
+    assert out_half.dtype == torch.float16
+
+    # bfloat16 may not be available on all GPUs; guard with try
+    try:
+        x_bf16 = torch.randn(2, 10, 16, device="cuda", dtype=torch.bfloat16)
+        out_bf16 = split_heads(x_bf16, 4)
+        assert out_bf16.dtype == torch.bfloat16
+    except RuntimeError:
+        pytest.skip("bfloat16 not supported on this device")
+
+
+###___calculate_attention___###
+
+
+# ----------------------------
+# Helpers
+# ----------------------------
+def _rand(B=2, H=3, Sq=4, Sk=5, D=8, device="cpu", dtype=torch.float32):
+    q = torch.randn(B, H, Sq, D, device=device, dtype=dtype)
+    k = torch.randn(B, H, Sk, D, device=device, dtype=dtype)
+    v = torch.randn(B, H, Sk, D, device=device, dtype=dtype)
+    return q, k, v
+
+
+# ----------------------------
+# A. Basic correctness
+# ----------------------------
+def test_basic_shapes_and_row_sums():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q, k, v = _rand(B, H, Sq, Sk, D)
+    attn, probs = calculate_attention(q, k, v, mask=None, return_probs=True)
+    assert attn.shape == (B, H, Sq, D)
+    assert probs.shape == (B, H, Sq, Sk)
+    # softmax rows sum to 1
+    assert torch.allclose(probs.sum(-1), torch.ones(B, H, Sq), atol=1e-6)
+
+
+def test_identity_prefers_diagonal():
+    # q == k == v as orthogonal-ish basis to encourage diagonal peak
+    Sq = Sk = D = 4
+    base = torch.eye(D).view(1, 1, Sk, D).expand(1, 1, Sk, D).contiguous()
+    q = base.clone()
+    k = base.clone()
+    v = base.clone()
+    _, probs = calculate_attention(q, k, v, mask=None, return_probs=True)
+    assert torch.equal(probs[0, 0].argmax(-1), torch.arange(Sq))
+
+
+# ----------------------------
+# B. Mask behavior (boolean & causal)
+# ----------------------------
+def test_padding_mask_boolean_last_two_keys():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q, k, v = _rand(B, H, Sq, Sk, D)
+    mask = torch.zeros(B, H, Sq, Sk, dtype=torch.bool)
+    mask[:, :, :, -2:] = True  # mask last two keys
+    _, probs = calculate_attention(q, k, v, mask=mask, return_probs=True)
+    assert (probs[:, :, :, -2:] < 1e-6).all()
+    # After zeroing masked entries, rows renormalize to ~1
+    assert torch.allclose(probs.masked_fill(mask, 0).sum(-1), torch.ones(B, H, Sq), atol=1e-6)
+
+
+def test_causal_mask_upper_triangle_zero():
+    B, H, S, D = 2, 3, 6, 8
+    q, k, v = _rand(B, H, S, S, D)
+    causal = torch.ones(B, H, S, S, dtype=torch.bool).triu(1)  # True above diagonal
+    _, probs = calculate_attention(q, k, v, mask=causal, return_probs=True)
+    assert (probs.triu(1) < 1e-6).all()
+
+
+def test_fully_masked_row_zero_probs_and_output():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q, k, v = _rand(B, H, Sq, Sk, D)
+    mask = torch.zeros(B, H, Sq, Sk, dtype=torch.bool)
+    mask[:, :, 1, :] = True  # fully mask a query row
+    attn, probs = calculate_attention(q, k, v, mask=mask, return_probs=True)
+    assert torch.allclose(probs[:, :, 1, :], torch.zeros(B, H, Sk))
+    assert torch.allclose(attn[:, :, 1, :], torch.zeros(B, H, D))
+
+
+# ----------------------------
+# C. Numerical stability
+# ----------------------------
+def test_extreme_logits_no_nans_or_infs():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q, k, v = _rand(B, H, Sq, Sk, D)
+    q = q * 1000
+    k = k * 1000
+    attn, probs = calculate_attention(q, k, v, mask=None, return_probs=True)
+    assert not (torch.isnan(probs).any() or torch.isinf(probs).any())
+    assert attn.shape == (B, H, Sq, D)
+
+
+def test_half_precision_close_to_fp32_or_skip():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q32, k32, v32 = _rand(B, H, Sq, Sk, D, device="cuda", dtype=torch.float32)
+    q16, k16, v16 = q32.half(), k32.half(), v32.half()
+    a16 = calculate_attention(q16, k16, v16, mask=None)
+    a32 = calculate_attention(q32, k32, v32, mask=None)
+    # looser tolerance for fp16
+    assert torch.allclose(a16.float(), a32.float(), atol=5e-2, rtol=5e-2)
+
+
+# ----------------------------
+# D. Edge cases
+# ----------------------------
+def test_zero_length_sequences():
+    B, H, D = 2, 3, 8
+    # Sq==0
+    q, k, v = _rand(B, H, 0, 5, D)
+    a, p = calculate_attention(q, k, v, mask=None, return_probs=True)
+    assert a.shape == (B, H, 0, D) and p.shape == (B, H, 0, 5)
+    # Sk==0
+    q, k, v = _rand(B, H, 4, 0, D)
+    a, p = calculate_attention(q, k, v, mask=None, return_probs=True)
+    assert a.shape == (B, H, 4, D) and p.shape == (B, H, 4, 0)
+    assert torch.allclose(a, torch.zeros_like(a))
+
+
+def test_singleton_softmax_is_one():
+    q = torch.randn(1, 1, 1, 1)
+    k = torch.randn(1, 1, 1, 1)
+    v = torch.randn(1, 1, 1, 1)
+    _, p = calculate_attention(q, k, v, mask=None, return_probs=True)
+    assert torch.allclose(p, torch.ones_like(p), atol=1e-6)
+
+
+# ----------------------------
+# E. Autograd & determinism
+# ----------------------------
+def test_gradients_flow_no_inplace_breakage():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q = torch.randn(B, H, Sq, D, requires_grad=True)
+    k = torch.randn(B, H, Sk, D, requires_grad=True)
+    v = torch.randn(B, H, Sk, D, requires_grad=True)
+    out = calculate_attention(q, k, v, mask=None)
+    loss = out.pow(2).sum()
+    loss.backward()
+    for t in (q, k, v):
+        assert t.grad is not None and torch.isfinite(t.grad).all()
+
+
+def test_deterministic_best_effort_same_seed():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q, k, v = _rand(B, H, Sq, Sk, D)
+    torch.manual_seed(42)
+    a1 = calculate_attention(q, k, v, mask=None, deterministic=True)
+    torch.manual_seed(42)
+    a2 = calculate_attention(q, k, v, mask=None, deterministic=True)
+    assert torch.allclose(a1, a2)
+
+
+# ----------------------------
+# F. Parity with PyTorch SDPA (when available)
+# ----------------------------
+def test_sdpa_parity_if_available():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q, k, v = _rand(B, H, Sq, Sk, D)
+    try:
+        sdpa = F.scaled_dot_product_attention(
+            q, k, v, dropout_p=0.0, attn_mask=None, is_causal=False
+        )
+        manual = calculate_attention(q, k, v, mask=None)
+        assert torch.allclose(sdpa, manual, atol=1e-5, rtol=1e-4)
+    except Exception:
+        # Older PyTorch: skip
+        pass
+
+
+# ----------------------------
+# G. Performance smoke (sanity only)
+# ----------------------------
+def test_perf_smoke_runs_reasonably():
+    # Not asserting timing; just ensure no OOM or pathological slowdowns
+    q, k, v = _rand(B=1, H=8, Sq=1024, Sk=1024, D=64)
+    _ = calculate_attention(q, k, v, mask=None, return_probs=False)
+    q, k, v = _rand(B=2, H=8, Sq=1536, Sk=1536, D=64)
+    _ = calculate_attention(q, k, v, mask=None, return_probs=False)
+
+
+# ----------------------------
+# H. Device behavior
+# ----------------------------
+def test_mask_on_cpu_qkv_on_gpu_autofix_or_skip():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
+    prev = torch.are_deterministic_algorithms_enabled()
+    try:
+        # Disable determinism *only for this test*
+        if prev:
+            torch.use_deterministic_algorithms(False)
+
+        q, k, v = _rand(device="cuda")
+        mask = torch.zeros(q.shape[0], 1, 1, k.shape[2], dtype=torch.bool)  # CPU
+        attn, probs = calculate_attention(q, k, v, mask, return_probs=True)
+        assert attn.device.type == "cuda"
+        assert probs.device.type == "cuda"
+    finally:
+        # restore whatever the global setting was
+        torch.use_deterministic_algorithms(prev)
+
+
+def test_qkv_device_mismatch_raises_clear_error():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
+    prev = torch.are_deterministic_algorithms_enabled()
+    try:
+        if prev:
+            torch.use_deterministic_algorithms(False)
+
+        q, k, v = _rand(device="cuda")
+        k = k.cpu()  # force mismatch
+        with pytest.raises(RuntimeError) as ei:
+            _ = calculate_attention(q, k, v, mask=None)
+        msg = str(ei.value)
+        assert (
+            "q/k/v must be on the same device" in msg  # our explicit check
+            or "Expected all tensors to be on the same device" in msg  # PyTorch matmul
+        )
+    finally:
+        torch.use_deterministic_algorithms(prev)
+
+
+# ----------------------------
+# I. Mask broadcastability
+# ----------------------------
+def test_boolean_mask_broadcast_variants_equivalent():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q, k, v = _rand(B, H, Sq, Sk, D)
+    m_exact = torch.zeros(B, H, Sq, Sk, dtype=torch.bool)
+    m_B1_1Sk = torch.zeros(B, 1, 1, Sk, dtype=torch.bool)
+    m_11SqSk = torch.zeros(1, 1, Sq, Sk, dtype=torch.bool)
+
+    a_exact, p_exact = calculate_attention(q, k, v, m_exact, return_probs=True)
+    a1, p1 = calculate_attention(q, k, v, m_B1_1Sk, return_probs=True)
+    a2, p2 = calculate_attention(q, k, v, m_11SqSk, return_probs=True)
+
+    assert torch.allclose(a1, a_exact, atol=1e-6, rtol=1e-5)
+    assert torch.allclose(p1, p_exact, atol=1e-6, rtol=1e-5)
+    assert torch.allclose(a2, a_exact, atol=1e-6, rtol=1e-5)
+    assert torch.allclose(p2, p_exact, atol=1e-6, rtol=1e-5)
+
+
+def test_additive_mask_broadcast_variants_equivalent():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q, k, v = _rand(B, H, Sq, Sk, D)
+    m_exact = torch.zeros(B, H, Sq, Sk, dtype=torch.float32)
+    m_111Sk = torch.zeros(1, 1, 1, Sk, dtype=torch.float32)
+    m_B1Sq1 = torch.zeros(B, 1, Sq, 1, dtype=torch.float32)
+
+    a_exact, p_exact = calculate_attention(q, k, v, m_exact, return_probs=True)
+    a1, p1 = calculate_attention(q, k, v, m_111Sk, return_probs=True)
+    a2, p2 = calculate_attention(q, k, v, m_B1Sq1, return_probs=True)
+
+    assert torch.allclose(a1, a_exact, atol=1e-6, rtol=1e-5)
+    assert torch.allclose(p1, p_exact, atol=1e-6, rtol=1e-5)
+    assert torch.allclose(a2, a_exact, atol=1e-6, rtol=1e-5)
+    assert torch.allclose(p2, p_exact, atol=1e-6, rtol=1e-5)
+
+
+def test_non_broadcastable_mask_raises_value_error():
+    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
+    q, k, v = _rand(B, H, Sq, Sk, D)
+    bad = torch.zeros(
+        B, H, Sq, dtype=torch.bool
+    )  # missing last dim => not broadcastable to (B,H,Sq,Sk)
+    with pytest.raises(ValueError) as ei:
+        _ = calculate_attention(q, k, v, bad)
+    assert "not broadcastable" in str(ei.value)
+
+
+###___join_heads___###
+
+
+def test_join_heads_type_error():
+    with pytest.raises(TypeError):
+        join_heads([1, 2, 3])  # not a tensor
+
+
+def test_join_heads_dim_error():
+    with pytest.raises(ValueError):
+        join_heads(torch.randn(2, 3, 4))  # only 3D
+
+
+def test_join_heads_shape_values():
+    x = torch.randn(2, 0, 4, 5)  # zero heads
+    y = join_heads(x)
+    assert y.shape == (2, 4, 0)
+
+
+def test_join_heads_roundtrip():
+    B, H, T, Dh = 2, 3, 5, 7
+    x = torch.randn(B, H, T, Dh)
+    y = join_heads(x)
+    assert y.shape == (B, T, H * Dh)
+    # Check element mapping correctness
+    for b in range(B):
+        for h in range(H):
+            for t in range(T):
+                for d in range(Dh):
+                    assert torch.allclose(y[b, t, h * Dh + d], x[b, h, t, d])
+
+
+###___sinusoidal_positional_encoding___###
+
+
+@pytest.mark.parametrize("seq_len,dim", [(0, 1), (1, 1), (1, 2), (7, 4), (17, 33)])
+def test_shapes_and_dtypes(seq_len, dim):
+    # Try a few dtypes; fp16 on CPU can be quirky, so we gate it on CUDA.
+    dtypes = [torch.float32, torch.bfloat16]
+    if torch.cuda.is_available():
+        dtypes.append(torch.float16)
+
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    for dtype in dtypes:
+        pe = sinusoidal_positional_encoding(seq_len, dim, dtype=dtype, device=device)
+        assert pe.shape == (seq_len, dim)
+        assert pe.dtype == dtype
+        # Returned tensor should be a constant table (no gradients)
+        assert pe.requires_grad is False
+
+
+def test_values_match_reference_small_case():
+    # Small deterministic comparison against the textbook formula.
+    seq_len, dim = 4, 6
+    base = 10_000.0
+    pe = sinusoidal_positional_encoding(seq_len, dim, dtype=torch.float32, device="cpu", base=base)
+
+    ref = torch.zeros_like(pe)
+    # Fill sin/cos pairs explicitly
+    for pos in range(seq_len):
+        for i in range(0, dim // 2):
+            denom = base ** (2 * i / dim)
+            ref[pos, 2 * i] = math.sin(pos / denom)
+            ref[pos, 2 * i + 1] = math.cos(pos / denom)
+
+    # Compare only the complete sin/cos pairs (even columns and their cos twins)
+    paired = (dim // 2) * 2
+    assert torch.allclose(pe[:, :paired], ref[:, :paired], atol=1e-6, rtol=1e-6)
+
+
+@pytest.mark.parametrize("dim", [1, 3, 8])
+def test_odd_dim_last_column_zero(dim):
+    seq_len = 5
+    pe = sinusoidal_positional_encoding(seq_len, dim)
+    if dim % 2 == 1:
+        assert torch.allclose(
+            pe[:, -1],
+            torch.zeros(seq_len, dtype=pe.dtype, device=pe.device),
+        )
+
+
+def test_offset_windowing_equivalence():
+    # pe[pos] with offset should match a slice of a longer table without offset
+    seq_len, dim, offset = 8, 6, 5
+    base_long = sinusoidal_positional_encoding(seq_len + offset, dim)
+    win = sinusoidal_positional_encoding(seq_len, dim, offset=offset)
+    assert torch.allclose(win, base_long[offset : offset + seq_len])
+
+
+def test_determinism():
+    a = sinusoidal_positional_encoding(32, 64)
+    b = sinusoidal_positional_encoding(32, 64)
+    assert torch.allclose(a, b)
+
+
+def test_invalid_inputs():
+    with pytest.raises(ValueError):
+        sinusoidal_positional_encoding(-1, 8)
+    with pytest.raises(ValueError):
+        sinusoidal_positional_encoding(1, 0)
+
+
+@pytest.mark.parametrize("base", [1_000.0, 10_000.0, 100_000.0])
+def test_different_bases_change_values(base):
+    seq_len, dim = 6, 8
+    a = sinusoidal_positional_encoding(seq_len, dim, base=base)
+    b = sinusoidal_positional_encoding(seq_len, dim, base=base * 10)
+    # Different bases should not produce identical tables
+    assert not torch.allclose(a, b)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for strict fp16 checks")
+def test_cuda_half_precision_support():
+    seq_len, dim = 12, 16
+    pe = sinusoidal_positional_encoding(seq_len, dim, dtype=torch.float16, device="cuda")
+    assert pe.device.type == "cuda"
+    assert pe.dtype == torch.float16
+
+
+###___combine_masks___###
+
+
+def test_both_none_returns_none():
+    assert combine_masks(None, None) is None
+
+
+def test_one_none_returns_other():
+    m = torch.tensor([[True, False]])
+    assert torch.equal(combine_masks(m, None), m)
+    assert torch.equal(combine_masks(None, m), m)
+
+
+def test_or_combination():
+    m1 = torch.tensor([[True, False]])
+    m2 = torch.tensor([[False, True]])
+    expected = torch.tensor([[True, True]])
+    out = combine_masks(m1, m2)
+    assert torch.equal(out, expected)
+
+
+def test_shape_mismatch_error():
+    m1 = torch.ones(2, 2, dtype=torch.bool)
+    m2 = torch.ones(3, 3, dtype=torch.bool)
+    with pytest.raises(ValueError):
+        combine_masks(m1, m2)
+
+
+def test_broadcastable_masks():
+    m1 = torch.tensor([[True, False, True]])
+    m2 = torch.tensor([[False]])  # broadcast across
+    out = combine_masks(m1, m2)
+    assert torch.equal(out, m1 | m2)
+
+
+def test_dtype_conversion():
+    # Non-bool masks still should work if they are 0/1 ints
+    m1 = torch.tensor([[1, 0]], dtype=torch.int32)
+    m2 = torch.tensor([[0, 1]], dtype=torch.int32)
+    out = combine_masks(m1.bool(), m2.bool())
+    expected = torch.tensor([[True, True]])
+    assert torch.equal(out, expected)
+
+
+###___shift_right___###
+
+
+def test_basic_shift_single_batch():
+    labels = torch.tensor([[5, 6, 7]], dtype=torch.long)
+    out = shift_right(labels, bos_id=1, pad_id=0)
+    assert torch.equal(out, torch.tensor([[1, 5, 6]], dtype=torch.long))
+
+
+def test_basic_shift_vector_input():
+    labels = torch.tensor([10, 11, 12], dtype=torch.long)
+    out = shift_right(labels, bos_id=2, pad_id=0)
+    assert out.shape == (1, 3)
+    assert torch.equal(out, torch.tensor([[2, 10, 11]], dtype=torch.long))
+
+
+def test_batch_shift():
+    labels = torch.tensor([[5, 6, 7], [8, 9, 10]], dtype=torch.long)
+    out = shift_right(labels, bos_id=3, pad_id=0)
+    expected = torch.tensor([[3, 5, 6], [3, 8, 9]], dtype=torch.long)
+    assert torch.equal(out, expected)
+
+
+def test_ignores_are_padded_after_shift():
+    # -100 in labels must never appear in inputs; becomes PAD after shifting.
+    labels = torch.tensor([[5, -100, 7, -100]], dtype=torch.long)
+    out = shift_right(labels, bos_id=4, pad_id=99)
+    # positions 2 and 4 in inputs come from -100 -> PAD(99)
+    expected = torch.tensor([[4, 5, 99, 7]], dtype=torch.long)
+    assert torch.equal(out, expected)
+    assert (out == -100).sum() == 0
+
+
+def test_custom_pad_and_bos_ids():
+    labels = torch.tensor([[42, 43, -100, 45]], dtype=torch.long)
+    out = shift_right(labels, bos_id=7, pad_id=3)
+    expected = torch.tensor([[7, 42, 43, 3]], dtype=torch.long)
+    assert torch.equal(out, expected)
+
+
+def test_zero_length_raises():
+    with pytest.raises(ValueError):
+        shift_right(torch.zeros((2, 0), dtype=torch.long))
+
+
+def test_wrong_rank_raises():
+    with pytest.raises(ValueError):
+        shift_right(torch.zeros((2, 3, 4), dtype=torch.long))
+
+
+def test_dtype_device_preserved():
+    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+    labels = torch.tensor([[1, 2, 3]], device=device)
+    out = shift_right(labels, bos_id=9, pad_id=0)
+    assert out.device == device
+    assert out.dtype == torch.long
+    assert torch.equal(out, torch.tensor([[9, 1, 2]], device=device, dtype=torch.long))
+
+
+def test_no_inplace_on_labels():
+    labels = torch.tensor([[5, 6, -100]], dtype=torch.long)
+    labels_clone = labels.clone()
+    _ = shift_right(labels, bos_id=1, pad_id=0)
+    assert torch.equal(labels, labels_clone)
+
+
+###___sample_from_logits___###
+
+
+def test_greedy_equals_argmax():
+    torch.manual_seed(0)
+    logits = torch.randn(4, 7)
+    ids = sample_from_logits(logits)  # greedy
+    ref = torch.argmax(F.softmax(logits, dim=-1), dim=-1)
+    assert torch.equal(ids, ref)
+
+
+def test_temperature_effect_no_sampling():
+    # Greedy with temperature should leave argmax unchanged (monotonic transform)
+    logits = torch.tensor([[0.1, 1.2, 0.3]], dtype=torch.float32)
+    ids_t1 = sample_from_logits(logits, do_sample=False, temperature=1.0)
+    ids_t05 = sample_from_logits(logits, do_sample=False, temperature=0.5)
+    assert torch.equal(ids_t1, ids_t05)  # argmax unchanged
+
+
+def test_top_k_limits_candidates():
+    logits = torch.tensor([[1.0, 0.9, 0.1, -1.0]], dtype=torch.float32)
+    # with top_k=1 greedy must select the max (index 0)
+    ids = sample_from_logits(logits, do_sample=False, top_k=1)
+    assert torch.equal(ids, torch.tensor([0]))
+    # with sampling, and top_k=1, the only candidate is index 0
+    gen = torch.Generator().manual_seed(123)
+    ids_s = sample_from_logits(logits, do_sample=True, top_k=1, rng=gen)
+    assert torch.equal(ids_s, torch.tensor([0]))
+
+
+def test_top_p_nucleus_filters_tail():
+    logits = torch.tensor([[5.0, 4.0, 3.0, -5.0]], dtype=torch.float32)  # probs heavily skewed
+    ids = sample_from_logits(logits, do_sample=False, top_p=0.6)  # should keep a minimal head
+    assert ids.item() in (0, 1)  # greediest is 0 anyway; check no crash
+
+
+def test_min_tokens_to_keep_safety():
+    logits = torch.tensor([[0.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
+    # Even with very small top_p, we keep at least 1 token
+    ids = sample_from_logits(logits, do_sample=False, top_p=0.01, min_tokens_to_keep=1)
+    assert ids.numel() == 1
+
+
+def test_allow_and_deny_lists():
+    logits = torch.tensor([[0.1, 2.0, 0.3, 3.0]], dtype=torch.float32)  # best is idx=3
+    # Allow only {0,2} → best among those is 2
+    ids = sample_from_logits(logits, do_sample=False, allowed_tokens=[0, 2])
+    assert torch.equal(ids, torch.tensor([2]))
+    # Deny {3} → next best is 1
+    ids = sample_from_logits(logits, do_sample=False, disallowed_tokens=[3])
+    assert torch.equal(ids, torch.tensor([1]))
+
+
+def test_rng_reproducibility_sampling():
+    logits = torch.tensor([[1.0, 1.0, 1.0, 1.0]], dtype=torch.float32)  # uniform
+    g1 = torch.Generator().manual_seed(42)
+    g2 = torch.Generator().manual_seed(42)
+    s1 = sample_from_logits(logits, do_sample=True, rng=g1)
+    s2 = sample_from_logits(logits, do_sample=True, rng=g2)
+    assert torch.equal(s1, s2)
+
+
+def test_device_is_preserved_cpu_cuda():
+    target_device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+    logits = torch.randn(3, 5, device=target_device)
+    out = sample_from_logits(logits)  # greedy by default
+    assert out.device == target_device
+    assert out.dtype == torch.long
+
+
+def test_handles_higher_rank_logits_flattening():
+    # [B, T, V] -> flattened internally; here we just ensure no crash and correct shape
+    logits = torch.randn(2, 3, 7)
+    out = sample_from_logits(logits)  # returns [B*T]
+    assert out.shape == (2 * 3,)
+
+
+###___create_causal_mask___###
+
+
+def test_mask_shape_and_dtype_cpu():
+    L = 5
+    mask = create_causal_mask(L)
+    assert mask.shape == (1, 1, L, L)
+    assert mask.dtype == torch.bool
+    assert mask.device.type == "cpu"
+
+
+def test_mask_upper_triangle():
+    L = 4
+    mask = create_causal_mask(L)
+    # For causal mask: mask[i,j] == True if j > i
+    mat = mask[0, 0].to(torch.int)
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1],
+            [0, 0, 1, 1],
+            [0, 0, 0, 1],
+            [0, 0, 0, 0],
+        ]
+    )
+    assert torch.equal(mat, expected)
+
+
+@pytest.mark.parametrize("dtype", [torch.bool, torch.float32])
+def test_mask_dtype(dtype):
+    L = 3
+    mask = create_causal_mask(L, dtype=dtype)
+    assert mask.dtype == dtype
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_device_cuda():
+    L = 6
+    mask = create_causal_mask(L, device=torch.device("cuda"))
+    assert mask.device.type == "cuda"
+
+
+def test_invalid_length_raises():
+    with pytest.raises(ValueError):
+        create_causal_mask(0)