import unittest from types import SimpleNamespace import torch import torch.nn as nn from models.vision_embedder import VisionEmbedder # --- Test fixtures ----------------------------------------------------------- def make_config(**overrides): """ Build a minimal config object exposing exactly the attributes VisionEmbedder reads. Using a SimpleNamespace (rather than the full VLMConfig) keeps the tests self-contained and lets each test override a single dimension in isolation. Parameters: * overrides : keyword overrides applied on top of the small defaults; requires each key to name a real VisionEmbedder config attribute (model_flat_patch_dim, emb_out_dim, emb_ln_eps, posemb_table_size, pos_embd_table_initializer_range) Returns: A SimpleNamespace carrying the five config fields, small enough for hand-checkable tests (defaults: flat patch dim 12, output dim 8, table size 10). """ defaults = dict( model_flat_patch_dim=12, emb_out_dim=8, emb_ln_eps=1e-5, posemb_table_size=10, pos_embd_table_initializer_range=0.02, ) defaults.update(overrides) return SimpleNamespace(**defaults) def make_embedder(seed=0, **overrides): """ Construct a seeded VisionEmbedder so parameter initialization is reproducible. Parameters: * seed (int) : seed applied to the global RNG immediately before construction; requires seed >= 0 * overrides : config overrides forwarded to make_config Returns: A VisionEmbedder in its freshly initialized state (LayerNorm affine params at weight=1/bias=0, pos table normally initialized). """ torch.manual_seed(seed) return VisionEmbedder(make_config(**overrides)) def make_linear_embedder(seed=0, **overrides): """ A VisionEmbedder whose three LayerNorms are swapped for identities. With the norms removed the forward pass collapses to the affine, fully inspectable map x -> fc(x) + pos(positions), which lets tests assert the positional-embedding contribution exactly without fighting LayerNorm's nonlinearity. Parameters: * seed (int) : seed forwarded to make_embedder; requires seed >= 0 * overrides : config overrides forwarded to make_config Returns: A VisionEmbedder with ln1, ln2, ln3 replaced by nn.Identity(). """ embedder = make_embedder(seed, **overrides) embedder.ln1 = nn.Identity() embedder.ln2 = nn.Identity() embedder.ln3 = nn.Identity() return embedder def random_positions(num_patches, table_size, seed=0): """ Draw valid (num_patches, 2) integer XY positions inside the table's index range. Parameters: * num_patches (int) : number of patch rows to generate; requires num_patches >= 1 * table_size (int) : exclusive upper bound for each coordinate; requires table_size >= 1 so coordinates land in [0, table_size) * seed (int) : seed applied before sampling; requires seed >= 0 Returns: A (num_patches, 2) int64 tensor of coordinates, each in [0, table_size). """ generator = torch.Generator().manual_seed(seed) return torch.randint(0, table_size, (num_patches, 2), generator=generator) def reference_pos(pos_table, positions): """ Independent, loop-based oracle for the factorized positional embedding, used so the tests do not merely restate the vectorized `pos_table[positions, axes]` gather. Parameters: * pos_table (torch.Tensor) : the (table_size, 2, D) positional table; requires axis 1 to be the axis selector with index 0 == x and index 1 == y * positions (torch.Tensor) : a (num_patches, 2) integer tensor of [x, y] coordinates; requires every coordinate to be in [0, table_size) Returns: A (num_patches, D) tensor whose row i equals pos_table[x_i, 0] + pos_table[y_i, 1] — the x-axis embedding plus the y-axis embedding for patch i. """ num_patches = positions.shape[0] embed_dim = pos_table.shape[-1] out = torch.zeros(num_patches, embed_dim, dtype=pos_table.dtype) for i in range(num_patches): x = int(positions[i, 0]) y = int(positions[i, 1]) out[i] = pos_table[x, 0] + pos_table[y, 1] return out # --- Shape and dtype --------------------------------------------------------- class TestVisionEmbedderShapesAndDtype(unittest.TestCase): """Behavioral spec for VisionEmbedder output shape and dtype handling.""" def test_output_shape(self): # (N, model_flat_patch_dim) patches -> (N, emb_out_dim) embeddings. embedder = make_embedder() x = torch.randn(5, 12) positions = random_positions(5, 10) out = embedder(x, positions) self.assertEqual(out.shape, (5, 8)) def test_single_patch(self): # A degenerate batch of one patch must still produce one embedding row. embedder = make_embedder() out = embedder(torch.randn(1, 12), random_positions(1, 10)) self.assertEqual(out.shape, (1, 8)) def test_output_dtype_follows_weights_not_input(self): # Step 0 casts the input to the projection weight dtype, so a float32-weighted # module returns float32 even when handed a float64 input. embedder = make_embedder() # parameters are float32 by default out = embedder(torch.randn(4, 12, dtype=torch.float64), random_positions(4, 10)) self.assertEqual(out.dtype, torch.float32) def test_runs_in_double_precision(self): # When the whole module is float64, a float32 input is up-cast and the # output is float64 — i.e. the dtype tracks the weights, in either direction. embedder = make_embedder().double() out = embedder(torch.randn(4, 12, dtype=torch.float32), random_positions(4, 10)) self.assertEqual(out.dtype, torch.float64) def test_float_positions_are_rejected(self): # Positions index the table via advanced indexing, which is only defined for # integer tensors; a float positions tensor must raise rather than silently # round or truncate. embedder = make_embedder() float_positions = random_positions(4, 10).to(torch.float32) with self.assertRaises((IndexError, RuntimeError)): embedder(torch.randn(4, 12), float_positions) # --- Padding assertion ------------------------------------------------------- class TestVisionEmbedderPaddingAssertion(unittest.TestCase): """Behavioral spec for the no-(-1)-padding precondition.""" def test_raises_when_all_positions_negative(self): # The module's design assumes patches arrive already flattened with no padding; # an all -1 positions tensor must trip the guard. embedder = make_embedder() positions = torch.full((3, 2), -1) with self.assertRaises(AssertionError): embedder(torch.randn(3, 12), positions) def test_raises_on_a_single_negative_entry(self): # Even one -1 (a single padding coordinate) violates the precondition: the # guard is `.all()`, not "mostly". embedder = make_embedder() positions = random_positions(4, 10) positions[2, 1] = -1 with self.assertRaises(AssertionError): embedder(torch.randn(4, 12), positions) def test_allows_all_zero_positions(self): # Zero is a valid coordinate (top-left patch), so an all-zeros positions tensor # must pass the guard and produce finite output. embedder = make_embedder() out = embedder(torch.randn(4, 12), torch.zeros(4, 2, dtype=torch.long)) self.assertEqual(out.shape, (4, 8)) self.assertTrue(torch.isfinite(out).all()) # --- Positional embedding semantics ----------------------------------------- class TestVisionEmbedderPositionalEmbedding(unittest.TestCase): """Behavioral spec for the factorized 2D positional embedding.""" def test_matches_loop_oracle_with_identity_norms(self): # With the norms removed the forward is exactly fc(x) + pos; cross-check the # positional term against the independent loop oracle. embedder = make_linear_embedder() x = torch.randn(6, 12) positions = random_positions(6, 10) out = embedder(x, positions) expected = embedder.fc(x) + reference_pos(embedder.pos_embd_table.detach(), positions) self.assertTrue(torch.allclose(out, expected, atol=1e-6)) def test_axis_zero_is_x_axis_one_is_y(self): # Pin the axis convention: with a hand-built table and zero input, the output # must be table[x, 0] + table[y, 1]. Swapping the coordinate proves x and y read # *different* axis slices. embedder = make_linear_embedder() embed_dim = embedder.emb_out_dim with torch.no_grad(): embedder.pos_embd_table.zero_() embedder.pos_embd_table[2, 0] = torch.arange(embed_dim, dtype=torch.float32) # x-axis, coord 2 embedder.pos_embd_table[3, 1] = torch.full((embed_dim,), 5.0) # y-axis, coord 3 zero_input = torch.zeros(1, embedder.model_flat_patch_dim) # fc still contributes its bias on a zero input; isolate the positional term # by subtracting that baseline (fc(0) == bias). baseline = embedder.fc(zero_input)[0].detach() # position (x=2, y=3) hits both populated slots out_hit = embedder(zero_input, torch.tensor([[2, 3]])) self.assertTrue(torch.allclose( out_hit[0] - baseline, torch.arange(embed_dim, dtype=torch.float32) + 5.0, atol=1e-6)) # position (x=3, y=2) reads table[3,0] and table[2,1], both still zero out_miss = embedder(zero_input, torch.tensor([[3, 2]])) self.assertTrue(torch.allclose(out_miss[0] - baseline, torch.zeros(embed_dim), atol=1e-6)) def test_factorized_decomposition_x_and_y_separable(self): # Because pos(x,y) = table[x,0] + table[y,1], holding x fixed and changing y must # shift the embedding by exactly table[y,1]-table[y',1] (the x term cancels), and # symmetrically for the y-fixed case. embedder = make_linear_embedder() zero = torch.zeros(1, embedder.model_flat_patch_dim) table = embedder.pos_embd_table.detach() same_x = embedder(zero, torch.tensor([[2, 3]]))[0] - embedder(zero, torch.tensor([[2, 7]]))[0] self.assertTrue(torch.allclose(same_x, table[3, 1] - table[7, 1], atol=1e-6)) same_y = embedder(zero, torch.tensor([[2, 5]]))[0] - embedder(zero, torch.tensor([[6, 5]]))[0] self.assertTrue(torch.allclose(same_y, table[2, 0] - table[6, 0], atol=1e-6)) def test_identical_positions_give_identical_positional_term(self): # Two patches at the same grid coordinate must receive the same positional # contribution. Feeding identical inputs, their full outputs must match exactly. embedder = make_linear_embedder() x_row = torch.randn(1, 12) x = torch.cat([x_row, x_row], dim=0) positions = torch.tensor([[4, 1], [4, 1]]) out = embedder(x, positions) self.assertTrue(torch.allclose(out[0], out[1], atol=1e-6)) def test_position_is_not_x_y_symmetric(self): # pos(a,b) = table[a,0]+table[b,1] differs from pos(b,a) = table[b,0]+table[a,1] # whenever the axis tables differ, confirming the two axes are not interchangeable. embedder = make_linear_embedder() zero = torch.zeros(1, embedder.model_flat_patch_dim) forward = embedder(zero, torch.tensor([[2, 6]])) swapped = embedder(zero, torch.tensor([[6, 2]])) self.assertFalse(torch.allclose(forward, swapped, atol=1e-4)) # --- Per-patch independence (the flat-concatenation contract) ---------------- class TestVisionEmbedderPerPatchIndependence(unittest.TestCase): """ Behavioral spec for row independence — the property that justifies the flat (sum_i num_model_patches_i, D) layout instead of a padded (batch, num_patches, D) one. Every op acts per row, so patches never influence each other. """ def test_permutation_equivariance(self): # Reordering patches (and their positions identically) must reorder the outputs # the same way: f(x[p])[i] == f(x)[p][i]. embedder = make_embedder() x = torch.randn(7, 12) positions = random_positions(7, 10) perm = torch.randperm(7) out = embedder(x, positions) out_perm = embedder(x[perm], positions[perm]) self.assertTrue(torch.allclose(out_perm, out[perm], atol=1e-6)) def test_changing_one_patch_leaves_others_unchanged(self): # Mutating a single input row must not perturb any other output row — this is # what distinguishes the per-row LayerNorm here from a cross-row BatchNorm. embedder = make_embedder() x = torch.randn(5, 12) positions = random_positions(5, 10) out = embedder(x, positions) x2 = x.clone() x2[2] = torch.randn(12) out2 = embedder(x2, positions) untouched = [i for i in range(5) if i != 2] self.assertTrue(torch.allclose(out[untouched], out2[untouched], atol=1e-6)) self.assertFalse(torch.allclose(out[2], out2[2], atol=1e-6)) def test_concatenation_equivalence(self): # Embedding two patch sets stacked together equals stacking their separate # embeddings. This is exactly why concatenating all images' patches into one # flat batch is safe. embedder = make_embedder() x_a, pos_a = torch.randn(3, 12), random_positions(3, 10, seed=1) x_b, pos_b = torch.randn(4, 12), random_positions(4, 10, seed=2) joint = embedder(torch.cat([x_a, x_b]), torch.cat([pos_a, pos_b])) separate = torch.cat([embedder(x_a, pos_a), embedder(x_b, pos_b)]) self.assertTrue(torch.allclose(joint, separate, atol=1e-6)) # --- Normalization behavior -------------------------------------------------- class TestVisionEmbedderNormalization(unittest.TestCase): """Behavioral spec for the LayerNorm steps.""" def test_output_rows_are_normalized(self): # A freshly built module has ln3 affine at weight=1/bias=0, so each output row # must be (population) standardized: mean ~ 0 and variance ~ 1 across features. embedder = make_embedder(emb_out_dim=64) out = embedder(torch.randn(16, 12), random_positions(16, 10)) self.assertTrue(torch.allclose(out.mean(dim=-1), torch.zeros(16), atol=1e-5)) self.assertTrue(torch.allclose(out.var(dim=-1, unbiased=False), torch.ones(16), atol=1e-3)) def test_input_affine_invariance(self): # ln1 normalizes each raw patch, so any positive per-row affine transform of the # input (a*x + b) is erased: the output depends only on the input's *shape*, not # its scale or offset. Positions are held fixed. embedder = make_embedder() x = torch.randn(5, 12) positions = random_positions(5, 10) out = embedder(x, positions) out_affine = embedder(3.0 * x + 7.0, positions) self.assertTrue(torch.allclose(out, out_affine, atol=1e-4)) # --- Parameters -------------------------------------------------------------- class TestVisionEmbedderParameters(unittest.TestCase): """Behavioral spec for the learned parameters.""" def test_pos_table_shape(self): # The positional table is (table_size, 2, emb_out_dim): one row per coordinate # value, a size-2 axis dimension (x, y), and an embedding vector. embedder = make_embedder(posemb_table_size=43, emb_out_dim=8) self.assertEqual(tuple(embedder.pos_embd_table.shape), (43, 2, 8)) def test_pos_table_is_normally_initialized(self): # The table must be initialized with normal(0, initializer_range), not left at # zeros; check empirically on a large table that mean ~ 0 and std ~ range. init_range = 0.05 embedder = make_embedder( posemb_table_size=400, emb_out_dim=64, pos_embd_table_initializer_range=init_range ) table = embedder.pos_embd_table.detach() self.assertFalse(torch.all(table == 0)) self.assertAlmostEqual(float(table.mean()), 0.0, delta=5e-3) self.assertAlmostEqual(float(table.std()), init_range, delta=init_range * 0.1) if __name__ == "__main__": unittest.main()