""" e8_utils.py — fully self-contained E8 lattice utilities for the LatticeMemory HF Space. No imports from liora_core. All math is inlined. """ from __future__ import annotations import math from itertools import combinations import torch import torch.nn.functional as F # --------------------------------------------------------------------------- # Core E8 lattice math # --------------------------------------------------------------------------- def _decode_d8(x: torch.Tensor) -> torch.Tensor: """Round x to the nearest D8 lattice point (even-sum integer coordinates).""" z = torch.round(x) parity = z.sum(dim=-1) % 2 diff = x - z worst = diff.abs().argmax(dim=-1) adj = torch.sign(diff) adj = torch.where(adj == 0, torch.ones_like(adj), adj) mask = torch.zeros_like(x) mask.scatter_(-1, worst.unsqueeze(-1), 1.0) z_fixed = z + adj * mask return torch.where(parity.unsqueeze(-1) == 1, z_fixed, z) def _e8_nearest(x: torch.Tensor) -> torch.Tensor: """Return the nearest E8 lattice point to x (batched, last dim = 8).""" z0 = _decode_d8(x) z1 = _decode_d8(x - 0.5) + 0.5 d0 = (x - z0).pow(2).sum(dim=-1, keepdim=True) d1 = (x - z1).pow(2).sum(dim=-1, keepdim=True) return torch.where(d0 <= d1, z0, z1) # --------------------------------------------------------------------------- # Shell-1 codebook (240 vectors, shape [240, 8]) # --------------------------------------------------------------------------- def _build_shell1_codebook(device: torch.device = torch.device("cpu")) -> torch.Tensor: """Build the 240-vector E8 shell-1 codebook. E8 shell-1 consists of: - 112 vectors of the form (±1, ±1, 0, 0, 0, 0, 0, 0) in all permutations - 128 vectors of the form (±½, ±½, …, ±½) with an even number of minus signs Total: 112 + 128 = 240 vectors. """ vecs: list[list[float]] = [] # ±1 in two coordinates, rest 0 for i, j in combinations(range(8), 2): for si in (1.0, -1.0): for sj in (1.0, -1.0): v = [0.0] * 8 v[i] = si v[j] = sj vecs.append(v) # (±½)^8 with even number of minus signs for mask in range(256): signs = [1.0 if (mask >> bit) & 1 == 0 else -1.0 for bit in range(8)] if signs.count(-1.0) % 2 == 0: vecs.append([s * 0.5 for s in signs]) codebook = torch.tensor(vecs, dtype=torch.float32, device=device) if codebook.shape != (240, 8): raise RuntimeError( f"expected E8 shell-1 codebook shape (240, 8), got {tuple(codebook.shape)}" ) return codebook # --------------------------------------------------------------------------- # RF-Snap: batch snap embeddings to E8 # --------------------------------------------------------------------------- @torch.no_grad() def nestquant_snap(embeddings: torch.Tensor) -> torch.Tensor: """Snap a batch of embeddings [B, D] (D divisible by 8) to E8 lattice points. Each 8-dim block is independently scaled, snapped to the nearest E8 point, and rescaled back. The operation is a no-op in the limit of small beta. Args: embeddings: float32 tensor of shape [B, D]. Returns: Snapped embeddings of shape [B, D]. """ if embeddings.dim() != 2: raise ValueError(f"nestquant_snap expects [B, D], got {tuple(embeddings.shape)}") B, D = embeddings.shape if D % 8 != 0: raise ValueError(f"D={D} must be divisible by 8") blocks = embeddings.float().reshape(B, D // 8, 8) # [B, n_blocks, 8] beta = blocks.norm(p=2, dim=-1).clamp_min(1e-8) / math.sqrt(2.0) # [B, n_blocks] snapped = _e8_nearest(blocks / beta.unsqueeze(-1)) # [B, n_blocks, 8] return (snapped * beta.unsqueeze(-1)).reshape(B, D) # --------------------------------------------------------------------------- # E8 address: encode a single embedding as a hex string # --------------------------------------------------------------------------- @torch.no_grad() def embedding_to_e8_address(embedding: torch.Tensor, device: torch.device = torch.device("cpu")) -> str: """Convert a single embedding vector to its E8 lattice address (hex string). Each 8-dim block is mapped to an index in [0, 239] (the closest shell-1 codebook vector after unit-normalisation), then packed as a byte. The resulting byte string is returned as a lowercase hex string. Expected address length: (D // 8) * 2 hex chars. For D=1024 → 256 chars. Args: embedding: 1-D float tensor of shape [D]. device: torch device for codebook (CPU by default). Returns: Lowercase hex string of length D // 4. """ if embedding.dim() != 1: raise ValueError(f"embedding_to_e8_address expects 1-D tensor, got {tuple(embedding.shape)}") D = embedding.numel() if D % 8 != 0: raise ValueError(f"D={D} must be divisible by 8") codebook = _build_shell1_codebook(device) # [240, 8] vector = embedding.float().to(device) blocks = vector.reshape(-1, 8) # [n_blocks, 8] beta = blocks.norm(p=2, dim=-1).clamp_min(1e-8) / math.sqrt(2.0) # [n_blocks] snapped = _e8_nearest(blocks / beta.unsqueeze(-1)) # [n_blocks, 8] # Unit-normalise snapped blocks and project onto codebook snapped_unit = snapped / snapped.norm(dim=-1, keepdim=True).clamp_min(1e-8) dots = snapped_unit @ codebook.T / math.sqrt(2.0) # [n_blocks, 240] indices = dots.argmax(dim=-1) # [n_blocks], values in [0,239] return bytes(indices.tolist()).hex() # --------------------------------------------------------------------------- # Index size comparison # --------------------------------------------------------------------------- def index_size_bytes(n_docs: int, d_model: int) -> dict[str, int]: """Return byte counts for different index formats. Keys: float32 — raw fp32 storage (4 bytes/float) int4 — 4-bit quantization (0.5 bytes/float) rfsnap — RF-Snap E8 (3 bytes per 8-dim block: 1 byte index + 2 bytes fp16 scale) Args: n_docs: number of document vectors. d_model: embedding dimension (must be divisible by 8). Returns: Dict with keys 'float32', 'int4', 'rfsnap' and integer byte values. """ if d_model % 8 != 0: raise ValueError(f"d_model={d_model} must be divisible by 8") n_blocks = d_model // 8 return { "float32": n_docs * d_model * 4, "int4": n_docs * d_model // 2, "rfsnap": n_docs * n_blocks * 3, }