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
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
//! A codebook: a vocabulary of atomic symbols stored as a tile-aligned matrix so
//! that *cleanup* — finding the symbol most similar to a noisy hypervector — is a
//! single matrix multiply on the systolic engine.
//!
//! The codebook holds its `count` symbols as the columns of a `dim × count`
//! [`PaddedTileLattice`]. Resolving a batch of `b` query hypervectors is then the
//! product `Q (b × dim) · C (dim × count) → S (b × count)`, after which each row's
//! `argmax` names the recovered symbol. That product is exactly what a TPU's
//! matrix unit exists to do, and it resolves the *entire* batch against the
//! *entire* vocabulary in one shot.
use crate::bf16::Bf16;
use crate::geometry::Geometry;
use crate::hyper::Hyper;
use crate::lattice::PaddedTileLattice;
/// A vocabulary of `count` deterministic bipolar hypervectors of dimension `dim`.
#[derive(Clone)]
pub struct Codebook {
dim: usize,
count: usize,
seed: u64,
geom: Geometry,
/// Shape `dim × count`: column `j` is the atom for symbol `j`.
matrix: PaddedTileLattice<f32>,
}
impl Codebook {
/// Build a codebook of `count` symbols, each a `dim`-dimensional bipolar atom
/// drawn deterministically from `seed`.
pub fn new(dim: usize, count: usize, seed: u64) -> Self {
Codebook::with_geometry(dim, count, seed, Geometry::TPU_V)
}
/// Build a codebook with an explicit tile geometry.
pub fn with_geometry(dim: usize, count: usize, seed: u64, geom: Geometry) -> Self {
let mut dense = vec![0.0f32; dim * count];
for j in 0..count {
let atom = Hyper::atom(dim, seed, j as u64);
let coords = atom.as_slice();
for (p, &value) in coords.iter().enumerate() {
dense[p * count + j] = value;
}
}
let matrix = PaddedTileLattice::from_dense(dim, count, &dense, geom)
.expect("dense buffer is exactly dim*count");
Codebook {
dim,
count,
seed,
geom,
matrix,
}
}
/// Hypervector dimensionality.
#[inline]
pub fn dim(&self) -> usize {
self.dim
}
/// Number of symbols in the vocabulary.
#[inline]
pub fn len(&self) -> usize {
self.count
}
/// True if the codebook has no symbols.
#[inline]
pub fn is_empty(&self) -> bool {
self.count == 0
}
/// The seed this codebook was generated from.
#[inline]
pub fn seed(&self) -> u64 {
self.seed
}
/// Borrow the `dim × count` symbol matrix (column `j` is symbol `j`).
#[inline]
pub fn matrix(&self) -> &PaddedTileLattice<f32> {
&self.matrix
}
/// Regenerate symbol `id` as a hypervector.
pub fn atom(&self, id: usize) -> Hyper {
Hyper::atom(self.dim, self.seed, id as u64)
}
/// Resolve one noisy hypervector to its nearest symbol, returning
/// `(symbol_id, score)`. This is the matrix-vector form of cleanup.
pub fn cleanup(&self, query: &Hyper) -> (usize, f32) {
let mut best = (0usize, f32::NEG_INFINITY);
let scores = self.scores_batch(core::slice::from_ref(query));
for (j, &s) in scores[0].iter().enumerate() {
if s > best.1 {
best = (j, s);
}
}
best
}
/// Resolve a batch of noisy hypervectors against the whole vocabulary with a
/// single systolic matmul, returning the winning `(symbol_id, score)` per query.
pub fn cleanup_batch(&self, queries: &[Hyper]) -> Vec<(usize, f32)> {
let scores = self.scores_batch(queries);
scores
.into_iter()
.map(|row| {
let mut best = (0usize, f32::NEG_INFINITY);
for (j, s) in row.into_iter().enumerate() {
if s > best.1 {
best = (j, s);
}
}
best
})
.collect()
}
/// The raw `b × count` similarity scores for a batch of queries, computed as
/// the matmul `Q · C` on the systolic engine.
pub fn scores_batch(&self, queries: &[Hyper]) -> Vec<Vec<f32>> {
let b = queries.len();
if b == 0 {
return Vec::new();
}
let mut qdense = vec![0.0f32; b * self.dim];
for (i, q) in queries.iter().enumerate() {
debug_assert_eq!(q.dim(), self.dim, "query dim must match codebook dim");
qdense[i * self.dim..(i + 1) * self.dim].copy_from_slice(q.as_slice());
}
let qlat = PaddedTileLattice::from_dense(b, self.dim, &qdense, self.geom)
.expect("query buffer is exactly b*dim");
// The one matmul: (b × dim) · (dim × count) → (b × count).
let product = qlat
.matmul(&self.matrix)
.expect("contraction dim and geometry match by construction");
let dense = product.to_dense();
dense
.chunks_exact(self.count)
.map(|row| row.to_vec())
.collect()
}
/// Cleanup run in `bf16`, the dtype a TPU matrix unit actually consumes.
///
/// Both operands are narrowed to `bf16` and the matmul accumulates in `f32`,
/// exactly as a TPU does. Returns the winning `(symbol_id, score)` per query.
/// Recall is slightly lower than the `f32` path near capacity — that is the
/// precision/throughput trade an int8/bf16 accelerator makes.
pub fn cleanup_batch_bf16(&self, queries: &[Hyper]) -> Vec<(usize, f32)> {
let b = queries.len();
if b == 0 {
return Vec::new();
}
let mut qdense = vec![Bf16::ZERO; b * self.dim];
for (i, q) in queries.iter().enumerate() {
for (p, &value) in q.as_slice().iter().enumerate() {
qdense[i * self.dim + p] = Bf16::from_f32(value);
}
}
let qlat = PaddedTileLattice::from_dense(b, self.dim, &qdense, self.geom)
.expect("query buffer is exactly b*dim");
let cb_bf16 = self.matrix.map(|x| Bf16::from_f32(*x));
let product = qlat
.matmul(&cb_bf16)
.expect("contraction dim and geometry match by construction");
product
.to_dense()
.chunks_exact(self.count)
.map(|row| {
let mut best = (0usize, f32::NEG_INFINITY);
for (j, s) in row.iter().enumerate() {
let s = s.to_f32();
if s > best.1 {
best = (j, s);
}
}
best
})
.collect()
}
/// Re-superpose the codebook by a weight per symbol: `X · w`, a `dim`-vector.
///
/// This is the second matmul of a cleanup projection: given similarity weights
/// (e.g. from [`Codebook::scores_batch`]), it rebuilds the weighted sum of
/// atoms. Together, `superpose(scores_batch(v))` is the projection `X Xᵀ v`
/// that a resonator network iterates.
pub fn superpose(&self, weights: &[f32]) -> Hyper {
assert_eq!(weights.len(), self.count, "one weight per symbol");
let wlat = PaddedTileLattice::from_dense(self.count, 1, weights, self.geom)
.expect("weight buffer is exactly count*1");
// (dim × count) · (count × 1) → (dim × 1).
let product = self
.matrix
.matmul(&wlat)
.expect("contraction dim and geometry match by construction");
Hyper::from_vec(product.to_dense())
}
}
impl core::fmt::Debug for Codebook {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(
f,
"Codebook {{ dim: {}, symbols: {} }}",
self.dim, self.count
)
}
}