tests/test_embedder.py

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()