File size: 8,483 Bytes
895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | """Runtime spatial Set-BPE vocabulary and lossless row serialization."""
from __future__ import annotations
import heapq
from collections.abc import Iterable, Sequence
from dataclasses import dataclass
import numpy as np
from .._validation import active_dimensions
from ..native_backend import native_kernel
from ..configuration import _integer
from ..native_serialization import NativeCacheMixin
__all__ = ["SpatialVocabulary", "aggregate_records"]
def aggregate_records(
codes: np.ndarray, bins: Sequence[int] | None = None
) -> tuple[np.ndarray, np.ndarray]:
"""Lossless micro-aggregation of identical quantized actions."""
values = np.asarray(codes, dtype=np.int64)
if values.ndim != 2:
raise ValueError("codes must have shape [U, D]")
if bins is not None:
dimensions = tuple(int(count) for count in bins)
if len(dimensions) != values.shape[1]:
raise ValueError("bins must have one entry per code coordinate")
counts = np.asarray(dimensions, dtype=np.int64)
valid = values >= 0
if np.any(values < -1) or np.any(values >= counts[None, :]):
raise ValueError("codes must contain -1 or an in-range cell index")
cardinality = int(np.prod(np.asarray(dimensions, dtype=object)))
if np.all(valid) and cardinality <= np.iinfo(np.int64).max:
keys = np.ravel_multi_index(values.T, dimensions)
unique_keys, counts = np.unique(keys, return_counts=True)
unique = np.stack(np.unravel_index(unique_keys, dimensions), axis=1)
return unique.astype(np.int64, copy=False), counts.astype(np.int64)
unique, counts = np.unique(values, axis=0, return_counts=True)
return unique, counts.astype(np.int64)
@dataclass
class SpatialVocabulary(NativeCacheMixin):
"""Atoms, merge rules and the decoder metadata they induce."""
bins: tuple[int, ...]
merges: tuple[tuple[int, int], ...]
def __post_init__(self) -> None:
if not self.bins or any(count < 1 for count in self.bins):
raise ValueError("every dimension needs at least one bin")
self.offsets = np.concatenate(([0], np.cumsum(np.asarray(self.bins, dtype=np.int64))))
self.atom_count = int(self.offsets[-1])
self.vocab_size = self.atom_count + len(self.merges)
dimension = len(self.bins)
support = np.zeros(self.vocab_size, dtype=np.int64)
expansion = np.full((self.vocab_size, dimension), -1, dtype=np.int64)
for index, count in enumerate(self.bins):
for value in range(count):
token = int(self.offsets[index]) + value
support[token] = 1 << index
expansion[token, index] = value
for rank, (left, right) in enumerate(self.merges):
left = _integer(left, "spatial merge left")
right = _integer(right, "spatial merge right")
token = self.atom_count + rank
if left >= token or right >= token:
raise ValueError("a merge rule must reference already-created tokens")
if support[left] & support[right]:
raise ValueError("a merge rule joins two overlapping supports")
support[token] = support[left] | support[right]
merged = np.maximum(expansion[left], expansion[right])
expansion[token] = merged
self.support = support
self.expansion = expansion
self.anchor = np.asarray(
[int(mask & -mask).bit_length() - 1 for mask in support], dtype=np.int64
)
self.full_support = (1 << dimension) - 1
self._rules: dict[tuple[int, int], tuple[int, int]] = {}
for rank, (left, right) in enumerate(self.merges):
key = (min(left, right), max(left, right))
self._rules[key] = (self.atom_count + rank, rank)
@property
def dimension(self) -> int:
return len(self.bins)
def atom(self, dimension: int, value: int) -> int:
return int(self.offsets[dimension]) + int(value)
def encode(
self,
code: Sequence[int],
*,
active_dims: Sequence[int] | None = None,
) -> list[int]:
"""Deterministic online Set-BPE encoding (Algorithm 8)."""
compact = [int(value) for value in code]
if active_dims is None:
if len(compact) != self.dimension:
raise ValueError(f"a dense quantized action must have {self.dimension} slots")
pairs = enumerate(compact)
else:
dimensions = active_dimensions(active_dims, self.dimension)
if len(dimensions) != len(compact):
raise ValueError("active_dims must have one entry per compact coordinate")
pairs = zip(dimensions, compact, strict=True)
tokens = set()
for index, value in pairs:
if value == -1:
if active_dims is not None:
raise ValueError("compact coordinates must not contain -1")
continue
if not 0 <= value < self.bins[index]:
raise ValueError(f"bin {value} is out of range for dimension {index}")
tokens.add(self.atom(index, value))
if not tokens:
raise ValueError("a quantized action must contain at least one active coordinate")
constructor = native_kernel("SetBPE")
if constructor is not None:
codec = getattr(self, "_compiled_spatial", None)
if codec is None:
codec = self._compiled_spatial = constructor(self.bins, 1, self.merges, ())
dense = np.full((1, 1, self.dimension), -1, dtype=np.int64)
for atom in tokens:
dimension = int(self.anchor[atom])
dense[0, 0, dimension] = atom - self.offsets[dimension]
return codec.encode_batch(dense)[0]
heap: list[tuple[int, int, int, int]] = []
ordered = sorted(tokens)
for position, left in enumerate(ordered):
for right in ordered[position + 1 :]:
rule = self._rules.get((left, right))
if rule is not None:
heapq.heappush(heap, (rule[1], rule[0], left, right))
while heap:
_, child, left, right = heapq.heappop(heap)
if left not in tokens or right not in tokens:
continue
tokens.discard(left)
tokens.discard(right)
for other in tokens:
key = (min(child, other), max(child, other))
rule = self._rules.get(key)
if rule is not None:
heapq.heappush(heap, (rule[1], rule[0], key[0], key[1]))
tokens.add(child)
return self.serialize(tokens)
def serialize(self, tokens: Iterable[int]) -> list[int]:
"""Canonical order for autoregressive supervision."""
return sorted(tokens, key=lambda token: (int(self.anchor[token]), int(token)))
def decode(
self,
tokens: Sequence[int],
*,
active_dims: Sequence[int] | None = None,
) -> np.ndarray:
"""Expand tokens to a dense row, using ``-1`` for absent coordinates."""
kernel = native_kernel("spatial_expand")
if kernel is not None:
dims = active_dimensions(active_dims, self.dimension)
return kernel(
[int(token) for token in tokens],
self.expansion,
self.support,
sum(1 << d for d in dims),
)
code = np.full(self.dimension, -1, dtype=np.int64)
covered = 0
for token in tokens:
value = int(token)
if not 0 <= value < self.vocab_size:
raise ValueError(f"token {value} is outside the vocabulary")
if covered & int(self.support[value]):
raise ValueError("token supports overlap; this is not a valid action")
covered |= int(self.support[value])
local = self.expansion[value]
code = np.where(local >= 0, local, code)
if active_dims is None:
expected = self.full_support
else:
dimensions = active_dimensions(active_dims, self.dimension)
expected = sum(1 << dimension for dimension in dimensions)
if covered != expected:
raise ValueError("token set does not exactly cover the requested active dimensions")
return code
|