tests/units/modules/test_embedding.py

import math

import pytest
import torch

from mini_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_rejects_zero_length_max_seq():
    with pytest.raises(ValueError):
        InputEmbedding(11, 8, 0, 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) * (m.d_model**0.5)
    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 "Sequence length" in str(ei.value) and "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_shape():
    pe = PositionalEmbedding(4, 10)
    assert hasattr(pe, "pe")
    assert isinstance(pe.pe, torch.Tensor)
    assert pe.pe.shape == (1, 4, 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)