← Walkthroughmapa/models/pack_r4.py
Source snapshot · source SHA-256 a39de5062ec0
1"""Token coordinates and attention boundaries for a fixed session/window length.23Main path: build_r4_grid() prepares metadata; pack_band_tokens() arranges features.45Input: per-band tokens and contact metadata.6Output: flat tokens (B,total,d), matching coordinates, and array boundaries.7Next: towers.Encoder.forward_flat() adds region identity and runs the blocks.89Order: array -> contact -> band (slow, mid, fast) -> time within band.10Every contact has K = sum(T_b) tokens. Band tensors concatenate into11(B, C, K, d), then flatten to (B, C*K, d). R4Grid stores matching metadata,12not features; cu_seqlens marks each array's start and end in the token axis.13"""1415from __future__ import annotations1617from collections.abc import Sequence18from dataclasses import dataclass1920import torch21from torch import Tensor2223from mapa.models.geometry import L1Geometry24from mapa.models.stem import PER_BAND_SPECS2526# band axis order (matches stem.PER_BAND_SPECS and PerBandStem.forward): SLOW, MID, HGA.27BAND_STRIDES: tuple[int, ...] = tuple(st for _, st in PER_BAND_SPECS)28N_BANDS = len(PER_BAND_SPECS)293031@dataclass(frozen=True)32class R4Grid:33 """Clip-INDEPENDENT flat full-grid geometry (one token per valid (contact,band,pos))."""3435 depth: Tensor # (total,) long — clinical contact index per token (L1 index-RoPE)36 time_pos: Tensor # (total,) long — shared-32Hz lattice position (L1 time-RoPE)37 contact: Tensor # (total,) long — contact index into N (for gather/scatter)38 band: Tensor # (total,) long — 0=SLOW 1=MID 2=HGA39 bandpos: Tensor # (total,) long — band-token index j in [0, T_b)40 array: Tensor # (total,) long — array id per token41 cu_seqlens: Tensor # (S+1,) int32 — per-array block bounds, TOKEN units42 cu_seqlens_drop: Tensor # (S'+1,) int64 — cu minus zero-length arrays (njt)43 max_seqlen: int # longest array block, tokens (static bound)44 total: int # N * K_full — total tokens (constant per session)45 k_full: int # tokens per contact = sum_b T_b46 band_lengths: tuple[int, ...] # (T_slow, T_mid, T_hga)474849def band_token_counts(n_time: int) -> tuple[int, ...]:50 """(T_slow, T_mid, T_hga) = n_time // stride per band."""51 for st in BAND_STRIDES:52 if n_time % st != 0:53 raise ValueError(f"n_time {n_time} not a multiple of band stride {st}")54 return tuple(n_time // st for st in BAND_STRIDES)555657def _drop_offsets(cu: Tensor) -> Tensor:58 off = cu.to(torch.int64)59 keep = off[1:] != off[:-1]60 return torch.cat([off[:1], off[1:][keep]])616263def build_r4_grid(64 geom: L1Geometry, *, n_time: int, max_seqlen: int | None = None65) -> R4Grid:66 """Build the clip-independent flat full-grid layout for one session.6768 ``max_seqlen`` (longest array block, tokens) is a per-session CONSTANT; pass the cached69 value to skip its ``.item()`` host sync (the compiled per-step path). ``None`` ⇒ derive70 it here (one sync — the standalone/first-call fallback)."""71 device = geom.gather_idx.device72 S, max_c = geom.n_arrays, geom.max_c73 lengths = band_token_counts(n_time) # (T_slow, T_mid, T_hga)74 k_full = int(sum(lengths))7576 # canonical (array-major, slot) contact order.77 rows = torch.arange(S, device=device)[:, None].expand(S, max_c)[geom.valid] # (N,)78 canon = geom.gather_idx[geom.valid] # (N,) contact index into N79 depth_canon = geom.depth[geom.valid] # (N,) clinical depth80 n = int(canon.shape[0])8182 # the per-contact (band, bandpos, time_pos) block — identical for every contact.83 band_blk = torch.cat(84 [torch.full((L,), b, dtype=torch.long, device=device) for b, L in enumerate(lengths)]85 ) # (k_full,)86 pos_blk = torch.cat(87 [torch.arange(L, device=device) for L in lengths]88 ) # (k_full,) band-token index89 lat_blk = torch.cat(90 [torch.arange(L, device=device) * st for L, st in zip(lengths, BAND_STRIDES)]91 ) # (k_full,) lattice position = j * stride9293 # tile over the N contacts (contact-major ⇒ array-contiguous).94 depth = depth_canon.repeat_interleave(k_full)95 contact = canon.repeat_interleave(k_full)96 array = rows.repeat_interleave(k_full)97 band = band_blk.repeat(n)98 bandpos = pos_blk.repeat(n)99 time_pos = lat_blk.repeat(n)100101 # Per-array token counts = contacts * K. Two 4-contact arrays at K=104:102 # counts [416,416] -> boundaries [0,416,832].103 n_per_array = geom.valid.sum(dim=1).to(torch.long) # (S,)104 seg = (n_per_array * k_full).to(torch.int32) # (S,)105 cu = torch.zeros(S + 1, dtype=torch.int32, device=device)106 cu[1:] = seg.cumsum(0).to(torch.int32)107 if max_seqlen is None:108 max_seqlen = int((n_per_array.max() * k_full).item())109110 return R4Grid(111 depth=depth, time_pos=time_pos, contact=contact, band=band, bandpos=bandpos,112 array=array, cu_seqlens=cu, cu_seqlens_drop=_drop_offsets(cu),113 max_seqlen=max_seqlen, total=n * k_full, k_full=k_full, band_lengths=lengths,114 )115116117def pack_band_tokens(band_tokens: Sequence[Tensor], grid: R4Grid) -> Tensor:118 """Flatten per-band token tuples into the flat grid token order.119120 ``band_tokens`` is one tensor per band, each ``(B, n_contacts, n_tokens, d)``. The result121 is ``(B, total, d)`` in the order ``build_r4_grid`` lays out, so token ``t`` is the one122 whose contact, band and band position are ``grid.contact[t]``, ``grid.band[t]`` and123 ``grid.bandpos[t]``. The layout is not restated here: the contact order is read back off124 the grid, and the bands are concatenated in their fixed order.125 """126 if len(band_tokens) != len(grid.band_lengths):127 raise ValueError(128 f"expected {len(grid.band_lengths)} band token streams, got {len(band_tokens)}"129 )130 b0 = band_tokens[0]131 B, d = b0.shape[0], b0.shape[-1]132 n = grid.total // grid.k_full133 # 1. Read the array-grouped contact order from the metadata plan.134 canon = grid.contact.reshape(n, grid.k_full)[:, 0] # (n,) contact index, canonical order135 blocks: list[Tensor] = []136 for band, (tok, length) in enumerate(zip(band_tokens, grid.band_lengths)):137 if tok.shape[2] != length:138 raise ValueError(139 f"band {band} has {tok.shape[2]} tokens, grid expects {length}"140 )141 blocks.append(tok[:, canon]) # (B, n, T_b, d) reindexed to canonical contacts142 # 2. Concatenate each contact's bands, then flatten contacts and token positions.143 contact_tokens = torch.cat(blocks, dim=2) # (B,n,k_full,d): slow, mid, fast per contact144 flat_tokens = contact_tokens.reshape(B, grid.total, d) # (B,n*k_full,d)145 return flat_tokens