source_spatial.py

"""Question-conditioned recurrence over the complete native visual field.

The state is the dynamic-resolution sequence emitted by Qwen2.5-VL's vision
merger, before those image embeddings enter the language model.  Tokens keep
their native two-dimensional grid throughout the recurrence.  One shared cell
combines directional local messages with visual-to-question cross-attention::

    V_0 = VisionMerger(image)
    V_{k+1} = V_k + Cell(V_k, question, grid),  k = 0, ..., T - 1

Only ``V_T`` is inserted into the ordinary Qwen multimodal prefix.  The cell's
output projection is exactly zero at initialization, so every horizon,
including the default T=8, is initially bitwise identical to the base visual
embedding path.  The full language model is then run from scratch; no cache
created while constructing the visual field is available to answer decoding.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm


@dataclass
class SpatialVisualTelemetry:
    """Detached recurrence diagnostics; no visual trajectory is retained."""

    update_rms: torch.Tensor  # [T]
    update_relative: torch.Tensor  # [T]
    step_drift: torch.Tensor  # [T], 1 - cos(V_k, V_{k+1})
    final_relative: torch.Tensor  # scalar, RMS(V_T - V_0) / RMS(V_0)


class SpatialVisualRecurrentCell(nn.Module):
    """Shared O(N) spatial cell for packed, variable-resolution image tokens.

    Global N-by-N visual attention is deliberately absent: a sample may carry
    8192 merged visual tokens.  Four directional grid messages preserve local
    geometry in linear time, while low-rank visual-to-question attention makes
    every patch update task dependent.  The recurrent *state* remains at the
    native language-model width; ``inner_width`` limits only the update
    operator's compute and is not a state bottleneck.
    """

    def __init__(
        self,
        width: int,
        inner_width: int,
        num_heads: int,
        *,
        rms_norm_eps: float,
        residual_scale: float,
    ) -> None:
        super().__init__()
        if width < 1 or inner_width < 1:
            raise ValueError("width and inner_width must be positive")
        if num_heads < 1 or inner_width % num_heads:
            raise ValueError(
                "inner_width must be divisible by the positive num_heads"
            )
        if not 0.0 <= residual_scale <= 1.0:
            raise ValueError("residual_scale must lie in [0, 1]")
        self.width = int(width)
        self.inner_width = int(inner_width)
        self.num_heads = int(num_heads)
        self.head_dim = self.inner_width // self.num_heads
        self.residual_scale = float(residual_scale)

        self.state_norm = Qwen2RMSNorm(self.width, eps=rms_norm_eps)
        self.question_norm = Qwen2RMSNorm(self.width, eps=rms_norm_eps)

        # One projection contains distinct north/south/west/east maps.  This
        # keeps orientation identifiable without allocating a dense 3x3 dxd
        # convolution over the 3584-wide native visual state.
        self.center_proj = nn.Linear(self.width, self.inner_width, bias=False)
        self.neighbor_proj = nn.Linear(
            self.width, 4 * self.inner_width, bias=False
        )

        self.query_proj = nn.Linear(self.width, self.inner_width, bias=False)
        self.key_proj = nn.Linear(self.width, self.inner_width, bias=False)
        self.value_proj = nn.Linear(self.width, self.inner_width, bias=False)
        self.output_proj = nn.Linear(self.inner_width, self.width, bias=False)

    def zero_output(self) -> None:
        """Restore the exact base-model identity after generic HF init."""
        with torch.no_grad():
            self.output_proj.weight.zero_()

    @staticmethod
    def _token_layout(
        merged_grid_thw: torch.LongTensor,
        image_to_batch: torch.LongTensor,
        batch_size: int,
        expected_tokens: int,
    ) -> tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]:
        """Return token-to-batch ids, within-sample positions and counts."""
        if merged_grid_thw.ndim != 2 or merged_grid_thw.shape[-1] != 3:
            raise ValueError("merged_grid_thw must have shape [num_images, 3]")
        if image_to_batch.ndim != 1 or image_to_batch.shape[0] != len(
            merged_grid_thw
        ):
            raise ValueError("image_to_batch must contain one id per image")
        per_image = merged_grid_thw.prod(dim=-1).long()
        if int(per_image.sum()) != int(expected_tokens):
            raise ValueError(
                "merged grids and visual field disagree: "
                f"grid tokens={int(per_image.sum())}, state tokens={expected_tokens}"
            )
        token_to_batch = torch.repeat_interleave(image_to_batch, per_image)
        if token_to_batch.numel() and (
            int(token_to_batch.min()) < 0
            or int(token_to_batch.max()) >= batch_size
        ):
            raise ValueError("image_to_batch contains an out-of-range batch id")
        if token_to_batch.numel() > 1 and bool(
            (token_to_batch[1:] < token_to_batch[:-1]).any()
        ):
            raise ValueError(
                "images must follow the same batch-major order as Qwen inputs"
            )
        counts = torch.bincount(token_to_batch, minlength=batch_size)
        starts = counts.cumsum(0) - counts
        within = torch.arange(
            expected_tokens, device=token_to_batch.device
        ) - torch.repeat_interleave(starts, counts)
        return token_to_batch, within, counts

    @staticmethod
    def _neighbor_edges(
        merged_grid_thw: torch.LongTensor,
        expected_tokens: int,
    ) -> tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]:
        """Build packed directed 4-neighbour edges once for all recurrent steps.

        Direction ids are 0/1/2/3 = north/south/west/east *as seen by the
        destination token*.  Temporal slices are separate 2-D fields.
        """
        device = merged_grid_thw.device
        destinations: list[torch.Tensor] = []
        sources: list[torch.Tensor] = []
        directions: list[torch.Tensor] = []
        offset = 0
        for t_value, h_value, w_value in merged_grid_thw.detach().cpu().tolist():
            t, h, w = int(t_value), int(h_value), int(w_value)
            count = t * h * w
            if min(t, h, w) < 1:
                raise ValueError(f"invalid merged visual grid {(t, h, w)}")
            index = torch.arange(
                offset, offset + count, device=device, dtype=torch.long
            ).view(t, h, w)

            def append(dst: torch.Tensor, src: torch.Tensor, direction: int) -> None:
                if dst.numel() == 0:
                    return
                destinations.append(dst.reshape(-1))
                sources.append(src.reshape(-1))
                directions.append(
                    torch.full(
                        (dst.numel(),), direction, device=device, dtype=torch.long
                    )
                )

            append(index[:, 1:, :], index[:, :-1, :], 0)  # north
            append(index[:, :-1, :], index[:, 1:, :], 1)  # south
            append(index[:, :, 1:], index[:, :, :-1], 2)  # west
            append(index[:, :, :-1], index[:, :, 1:], 3)  # east
            offset += count
        if offset != expected_tokens:
            raise ValueError(
                f"edge grids contain {offset} tokens, expected {expected_tokens}"
            )
        if not destinations:
            empty = torch.empty(0, device=device, dtype=torch.long)
            return empty, empty, empty
        return (
            torch.cat(destinations),
            torch.cat(sources),
            torch.cat(directions),
        )

    def _question_attention(
        self,
        normalized_state: torch.Tensor,
        question_keys: torch.Tensor,
        question_values: torch.Tensor,
        question_valid: torch.BoolTensor,
        token_to_batch: torch.LongTensor,
        within_sample: torch.LongTensor,
        token_counts: torch.LongTensor,
    ) -> torch.Tensor:
        """Visual-query/text-memory SDPA for a packed visual sequence."""
        batch_size = question_keys.shape[0]
        max_tokens = int(token_counts.max()) if token_counts.numel() else 0
        if max_tokens == 0:
            return normalized_state.new_empty(0, self.inner_width)

        query_flat = self.query_proj(normalized_state)
        query = query_flat.new_zeros(
            batch_size, max_tokens, self.inner_width
        )
        query[token_to_batch, within_sample] = query_flat

        bsz, question_length, _ = question_keys.shape
        query = query.view(
            bsz, max_tokens, self.num_heads, self.head_dim
        ).transpose(1, 2)
        key = question_keys.view(
            bsz, question_length, self.num_heads, self.head_dim
        ).transpose(1, 2)
        value = question_values.view(
            bsz, question_length, self.num_heads, self.head_dim
        ).transpose(1, 2)
        # Boolean SDPA masks use True for entries that are allowed to attend.
        allowed = question_valid[:, None, None, :]
        attended = F.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=allowed,
            dropout_p=0.0,
            is_causal=False,
        )
        attended = attended.transpose(1, 2).reshape(
            bsz, max_tokens, self.inner_width
        )
        return attended[token_to_batch, within_sample]

    def forward(
        self,
        visual_state: torch.Tensor,
        merged_grid_thw: torch.LongTensor,
        image_to_batch: torch.LongTensor,
        question_embeddings: torch.Tensor,
        question_attention_mask: torch.Tensor | None,
        *,
        steps: int,
    ) -> tuple[torch.Tensor, SpatialVisualTelemetry]:
        if visual_state.ndim != 2 or visual_state.shape[-1] != self.width:
            raise ValueError(
                f"visual_state must be [N,{self.width}], got "
                f"{tuple(visual_state.shape)}"
            )
        if question_embeddings.ndim != 3 or question_embeddings.shape[-1] != self.width:
            raise ValueError(
                f"question_embeddings must be [B,Q,{self.width}]"
            )
        if steps < 0:
            raise ValueError("steps must be non-negative")
        batch_size = question_embeddings.shape[0]
        token_to_batch, within_sample, token_counts = self._token_layout(
            merged_grid_thw,
            image_to_batch,
            batch_size,
            visual_state.shape[0],
        )
        edge_dst, edge_src, edge_direction = self._neighbor_edges(
            merged_grid_thw, visual_state.shape[0]
        )

        if question_attention_mask is None:
            question_valid = torch.ones(
                question_embeddings.shape[:2],
                dtype=torch.bool,
                device=question_embeddings.device,
            )
        else:
            if question_attention_mask.shape != question_embeddings.shape[:2]:
                raise ValueError("question attention mask shape mismatch")
            question_valid = question_attention_mask > 0
        if not bool(question_valid.any(dim=-1).all()):
            raise ValueError("every image-bearing sample needs a question token")

        normalized_question = self.question_norm(question_embeddings)
        question_keys = self.key_proj(normalized_question)
        question_values = self.value_proj(normalized_question)

        initial = visual_state
        state = visual_state
        update_rms: list[torch.Tensor] = []
        update_relative: list[torch.Tensor] = []
        step_drift: list[torch.Tensor] = []
        for _ in range(steps):
            normalized = self.state_norm(state)
            center = self.center_proj(normalized)
            directional = self.neighbor_proj(normalized).view(
                state.shape[0], 4, self.inner_width
            )
            spatial = center.new_zeros(center.shape)
            degree = center.new_zeros(center.shape[0], 1)
            if edge_dst.numel():
                messages = directional[edge_src, edge_direction]
                spatial.index_add_(0, edge_dst, messages)
                degree.index_add_(
                    0,
                    edge_dst,
                    torch.ones(
                        edge_dst.shape[0], 1, dtype=center.dtype, device=center.device
                    ),
                )
                spatial = spatial / degree.clamp_min(1.0)
            question = self._question_attention(
                normalized,
                question_keys,
                question_values,
                question_valid,
                token_to_batch,
                within_sample,
                token_counts,
            )
            mixed = F.silu((center + spatial + question) / (3.0**0.5))
            update = self.output_proj(mixed) * self.residual_scale
            next_state = state + update

            with torch.no_grad():
                state_rms = state.float().square().mean().sqrt().clamp_min(1e-8)
                update_norm = update.float().square().mean().sqrt()
                cosine = F.cosine_similarity(
                    state.float(), next_state.float(), dim=-1
                ).mean()
                update_rms.append(update_norm.detach())
                update_relative.append((update_norm / state_rms).detach())
                step_drift.append((1.0 - cosine).detach())
            state = next_state

        with torch.no_grad():
            initial_rms = initial.float().square().mean().sqrt().clamp_min(1e-8)
            final_relative = (
                (state.float() - initial.float()).square().mean().sqrt()
                / initial_rms
            ).detach()
            empty = initial.new_empty(0, dtype=torch.float32)
            telemetry = SpatialVisualTelemetry(
                update_rms=(torch.stack(update_rms) if update_rms else empty),
                update_relative=(
                    torch.stack(update_relative) if update_relative else empty
                ),
                step_drift=(torch.stack(step_drift) if step_drift else empty),
                final_relative=final_relative,
            )
        return state, telemetry


__all__ = ["SpatialVisualRecurrentCell", "SpatialVisualTelemetry"]