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)