weftcodec 1.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
weftcodec/ai_modes.py ADDED
@@ -0,0 +1,3360 @@
1
+ """AI-mode payload decoders (WEFT-AI-CODEC §5 and §11-§17).
2
+
3
+ Modes implemented here:
4
+
5
+ * 9 ``FLOAT_TENSOR``
6
+ * 10 ``TENSOR_SHARDED``
7
+ * 11 ``WEIGHT_DELTA``
8
+ * 12 ``KV_CACHE``
9
+ * 13 ``ACTIVATION_STREAM``
10
+ * 14 ``SEMANTIC_DELTA``
11
+ * 15 ``TENSOR_TILED``
12
+ * 16 ``SPARSE_DELTA``
13
+ * 17 ``FLOAT_TILED``
14
+
15
+ Every function is an independent decoder written from the normative spec and
16
+ the canonical WCB reader in :mod:`weftcodec.wcb`. No Rust or Go code is loaded
17
+ or shelled out to. Error classes follow the normative validation order; there
18
+ is no approximate or lossy fallback.
19
+ """
20
+
21
+ from __future__ import annotations
22
+
23
+ import math
24
+ import struct
25
+ from dataclasses import dataclass
26
+
27
+ from . import huffman, varint, wcb
28
+ from .errors import (
29
+ Corrupt,
30
+ DependencyMismatch,
31
+ DependencyMissing,
32
+ LimitExceeded,
33
+ Malformed,
34
+ NonCanonical,
35
+ )
36
+ from .hashing import (
37
+ DOMAIN_AI_STREAM,
38
+ DOMAIN_TENSOR_LOGICAL,
39
+ DOMAIN_TENSOR_SEMANTIC,
40
+ DOMAIN_TENSOR_SEMANTIC_VALUES,
41
+ DOMAIN_TENSOR_TILE,
42
+ DOMAIN_TENSOR_FLOAT_TILE,
43
+ DOMAIN_TENSOR_FLOAT_TILE_STREAMS,
44
+ DOMAIN_TENSOR_VALUES,
45
+ DOMAIN_TENSOR_SHARDED_PAYLOAD,
46
+ hash_domain,
47
+ )
48
+ from .limits import Limits
49
+
50
+ U64_MAX = (1 << 64) - 1
51
+ U32_MAX = (1 << 32) - 1
52
+ HASH_LEN = 32
53
+
54
+ FLOAT_TENSOR_SCHEMA = "weft.ai.float_tensor.v1"
55
+ FLOAT_TENSOR_CODER = "huffman-canonical-v1"
56
+ TENSOR_SHARDED_SCHEMA = "weft.ai.tensor_sharded.v1"
57
+ WEIGHT_DELTA_SCHEMA = "weft.ai.weight_delta.v1"
58
+ SEMANTIC_DELTA_SCHEMA = "weft.ai.semantic_delta.v1"
59
+ KV_CACHE_SCHEMA = "weft.ai.kv_cache.v1"
60
+ ACTIVATION_STREAM_SCHEMA = "weft.ai.activation_stream.v1"
61
+ TENSOR_TILED_SCHEMA = "weft.ai.tensor_tiled.v1"
62
+ SPARSE_DELTA_SCHEMA = "weft.ai.sparse_delta.v1"
63
+ SEMANTIC_SCHEMA = "weft.tensor.semantic.v1"
64
+ LOGICAL_SCHEMA = "weft.tensor.logical.v1"
65
+ PLANE_CODEC = "huffman-canonical-v1"
66
+ DELTA_OP = "mod-sub"
67
+ SPARSE_CODEC = "zero-run-v1"
68
+ TILE_LAYOUT = "row_major"
69
+ BLOCK_SIZES = (16, 32, 64, 128, 256)
70
+ MAX_BLOCK_LEVELS = 1 << 24
71
+
72
+ FLOAT_FIELD_BITS = {
73
+ "bf16": (1, 8, 7),
74
+ "fp16": (1, 5, 10),
75
+ "fp32": (1, 8, 23),
76
+ "fp8-e4m3": (1, 4, 3),
77
+ "fp8-e5m2": (1, 5, 2),
78
+ }
79
+
80
+ WEIGHT_WORD_BITS = {"bf16": 16, "fp16": 16, "fp32": 32, "q8_0": 8, "q4_0": 4}
81
+ WEIGHT_STORAGE_BYTES = {"bf16": 2, "fp16": 2, "fp32": 4, "q8_0": 1, "q4_0": 1}
82
+ WEIGHT_PLANES = {"bf16": 2, "fp16": 2, "fp32": 4, "q8_0": 1, "q4_0": 1}
83
+ WEIGHT_MANTISSA_BITS = {"bf16": 7, "fp16": 10, "fp32": 23}
84
+
85
+ SHARD_STORAGE_BYTES = {"q8_0": 1, "q4_0": 1}
86
+
87
+ KV_WORD_BYTES = {"bf16": 2, "fp16": 2, "fp32": 4}
88
+ KV_PLANES = {"bf16": 2, "fp16": 2, "fp32": 4}
89
+
90
+ TENSOR_DTYPE_SIZES = {
91
+ "f16": 2,
92
+ "bf16": 2,
93
+ "f32": 4,
94
+ "f64": 8,
95
+ "i8": 1,
96
+ "u8": 1,
97
+ "i16": 2,
98
+ "i32": 4,
99
+ "i64": 8,
100
+ }
101
+
102
+
103
+ # ---------------------------------------------------------------------------
104
+ # Small checked-arithmetic and metadata helpers
105
+ # ---------------------------------------------------------------------------
106
+
107
+
108
+ def _checked_add(left: int, right: int, what: str) -> int:
109
+ value = left + right
110
+ if value > U64_MAX:
111
+ raise LimitExceeded("%s overflows u64" % what)
112
+ return value
113
+
114
+
115
+ def _checked_mul(left: int, right: int, what: str) -> int:
116
+ value = left * right
117
+ if value > U64_MAX:
118
+ raise LimitExceeded("%s overflows u64" % what)
119
+ return value
120
+
121
+
122
+ def _check_len(length: int, maximum: int, what: str) -> None:
123
+ if length > maximum:
124
+ raise LimitExceeded("%s length %d exceeds limit %d" % (what, length, maximum))
125
+
126
+
127
+ def _map(value, what: str) -> dict:
128
+ if not isinstance(value, dict):
129
+ raise Malformed("%s metadata must be a WCB map" % what)
130
+ return value
131
+
132
+
133
+ def _exact_keys(map_value: dict, expected, what: str) -> None:
134
+ for key in map_value:
135
+ if key not in expected:
136
+ raise Malformed("%s metadata has unknown key `%s`" % (what, key))
137
+ for key in expected:
138
+ if key not in map_value:
139
+ raise Malformed("%s metadata is missing key `%s`" % (what, key))
140
+
141
+
142
+ def _text(map_value: dict, key: str) -> str:
143
+ value = map_value.get(key)
144
+ if not isinstance(value, str):
145
+ raise Malformed("metadata `%s` must be text" % key)
146
+ return value
147
+
148
+
149
+ def _u64(map_value: dict, key: str) -> int:
150
+ value = map_value.get(key)
151
+ if not isinstance(value, wcb.U64):
152
+ raise Malformed("metadata `%s` must be u64" % key)
153
+ return value.value
154
+
155
+
156
+ def _u32(map_value: dict, key: str) -> int:
157
+ value = _u64(map_value, key)
158
+ if value > U32_MAX:
159
+ raise Malformed("metadata `%s` = %d does not fit u32" % (key, value))
160
+ return value
161
+
162
+
163
+ def _i64(map_value: dict, key: str) -> int:
164
+ value = map_value.get(key)
165
+ if not isinstance(value, wcb.I64):
166
+ raise Malformed("metadata `%s` must be i64" % key)
167
+ return value.value
168
+
169
+
170
+ def _bytes_of(map_value: dict, key: str) -> bytes:
171
+ value = map_value.get(key)
172
+ if not isinstance(value, (bytes, bytearray)):
173
+ raise Malformed("metadata `%s` must be bytes" % key)
174
+ return bytes(value)
175
+
176
+
177
+ def _hash_field(map_value: dict, key: str) -> bytes:
178
+ value = _bytes_of(map_value, key)
179
+ if len(value) != HASH_LEN:
180
+ raise Malformed(
181
+ "metadata `%s` must be %d bytes, got %d" % (key, HASH_LEN, len(value))
182
+ )
183
+ return value
184
+
185
+
186
+ def _array_u64(map_value: dict, key: str) -> "list[int]":
187
+ value = map_value.get(key)
188
+ if not isinstance(value, list):
189
+ raise Malformed("metadata `%s` must be an array" % key)
190
+ out = []
191
+ for item in value:
192
+ if not isinstance(item, wcb.U64):
193
+ raise Malformed("metadata `%s` elements must be u64" % key)
194
+ out.append(item.value)
195
+ return out
196
+
197
+
198
+ def _shape(map_value: dict) -> "list[int]":
199
+ return _array_u64(map_value, "shape")
200
+
201
+
202
+ def _split_prefix(body: bytes) -> "tuple[bytes, bytes]":
203
+ """Split ``u32le(meta_len) || descriptor || payload`` without a length bias."""
204
+ if len(body) < 4:
205
+ raise Malformed("AI body is shorter than its 4-byte metadata length")
206
+ meta_len = int.from_bytes(body[0:4], "little")
207
+ meta_end = 4 + meta_len
208
+ if meta_end > len(body):
209
+ raise Malformed("AI metadata extends past the frame body")
210
+ return body[4:meta_end], body[meta_end:]
211
+
212
+
213
+ def _split_exact(body: bytes, decoded_len: int) -> "tuple[bytes, bytes]":
214
+ descriptor, payload = _split_prefix(body)
215
+ if len(payload) != decoded_len:
216
+ raise Malformed(
217
+ "AI payload length %d differs from decoded_len %d"
218
+ % (len(payload), decoded_len)
219
+ )
220
+ return descriptor, payload
221
+
222
+
223
+ def _decode_descriptor(descriptor_bytes: bytes, limits: Limits):
224
+ return wcb.decode(descriptor_bytes, limits)
225
+
226
+
227
+ def _u64_vec(values) -> "list[wcb.U64]":
228
+ return [wcb.U64(value) for value in values]
229
+
230
+
231
+ def _validate_shape(shape: "list[int]", limits: Limits, what: str) -> int:
232
+ if not shape:
233
+ raise Malformed("%s shape rank must be at least 1" % what)
234
+ if len(shape) > limits.max_tensor_rank:
235
+ raise LimitExceeded(
236
+ "%s rank %d exceeds limit %d"
237
+ % (what, len(shape), limits.max_tensor_rank)
238
+ )
239
+ product = 1
240
+ for dim in shape:
241
+ if dim == 0:
242
+ raise Malformed("%s shape dimensions must be greater than zero" % what)
243
+ product = _checked_mul(product, dim, "%s shape product" % what)
244
+ if product > limits.max_tensor_elements:
245
+ raise LimitExceeded(
246
+ "%s element count %d exceeds limit %d"
247
+ % (what, product, limits.max_tensor_elements)
248
+ )
249
+ return product
250
+
251
+
252
+ # ---------------------------------------------------------------------------
253
+ # Logical roots and semantic roots
254
+ # ---------------------------------------------------------------------------
255
+
256
+
257
+ def values_hash(data: bytes) -> bytes:
258
+ return hash_domain(DOMAIN_TENSOR_VALUES, data)
259
+
260
+
261
+ def tensor_logical_root(
262
+ dtype: str, shape: "list[int]", element_count: int, data: bytes, limits: Limits
263
+ ) -> bytes:
264
+ record = {
265
+ "schema": LOGICAL_SCHEMA,
266
+ "dtype": dtype,
267
+ "shape": _u64_vec(shape),
268
+ "element_count": wcb.U64(element_count),
269
+ "values_hash": values_hash(data),
270
+ }
271
+ return hash_domain(DOMAIN_TENSOR_LOGICAL, wcb.encode(record, limits))
272
+
273
+
274
+ def sharded_logical_root(
275
+ dtype: str,
276
+ shape: "list[int]",
277
+ words: "list[int]",
278
+ scales: "list[int]",
279
+ zero_points: "list[int]",
280
+ limits: Limits,
281
+ ) -> bytes:
282
+ element_count = _validate_shape(shape, limits, "TENSOR_SHARDED")
283
+ if len(words) != element_count:
284
+ raise Malformed(
285
+ "TENSOR_SHARDED tensor has %d words but shape requires %d"
286
+ % (len(words), element_count)
287
+ )
288
+ if len(scales) != len(zero_points):
289
+ raise Malformed("TENSOR_SHARDED scales and zero_points differ in length")
290
+ material = bytearray()
291
+ for word in words:
292
+ if dtype == "q4_0" and not (-8 <= word <= 7):
293
+ raise Malformed("q4_0 word out of signed 4-bit range")
294
+ material.append(word & 0xFF)
295
+ for scale, zero in zip(scales, zero_points):
296
+ material.extend(struct.pack("<I", scale))
297
+ material.extend(struct.pack("<q", zero))
298
+ record = {
299
+ "schema": LOGICAL_SCHEMA,
300
+ "dtype": dtype,
301
+ "shape": _u64_vec(shape),
302
+ "element_count": wcb.U64(element_count),
303
+ "values_hash": values_hash(bytes(material)),
304
+ }
305
+ return hash_domain(DOMAIN_TENSOR_LOGICAL, wcb.encode(record, limits))
306
+
307
+
308
+ def tile_root(dtype: str, dims: "list[int]", tile_bytes: bytes, limits: Limits) -> bytes:
309
+ record = {
310
+ "dtype": dtype,
311
+ "dims": _u64_vec(dims),
312
+ "element_len": wcb.U64(len(tile_bytes)),
313
+ "values_hash": values_hash(tile_bytes),
314
+ }
315
+ return hash_domain(DOMAIN_TENSOR_TILE, wcb.encode(record, limits))
316
+
317
+
318
+ @dataclass(frozen=True)
319
+ class SemanticPolicy:
320
+ scheme: str
321
+ bits: int = 0
322
+ block_size: int = 0
323
+ levels: int = 0
324
+
325
+ def to_wcb(self):
326
+ if self.scheme == "none":
327
+ return {"scheme": "none"}
328
+ if self.scheme == "mantissa_truncate":
329
+ return {"bits": wcb.U64(self.bits), "scheme": "mantissa_truncate"}
330
+ return {
331
+ "block_size": wcb.U64(self.block_size),
332
+ "levels": wcb.U64(self.levels),
333
+ "scheme": "block_quantize",
334
+ }
335
+
336
+ def validate(self, dtype: str, element_count: int) -> None:
337
+ if self.scheme == "none":
338
+ return
339
+ if self.scheme == "mantissa_truncate":
340
+ maximum = WEIGHT_MANTISSA_BITS.get(dtype)
341
+ if maximum is None:
342
+ raise Malformed(
343
+ "mantissa_truncate policy is not defined for dtype %s" % dtype
344
+ )
345
+ if self.bits == 0 or self.bits > maximum:
346
+ raise Malformed(
347
+ "mantissa_truncate bits %d is outside 1..=%d for %s"
348
+ % (self.bits, maximum, dtype)
349
+ )
350
+ return
351
+ if self.block_size == 0:
352
+ raise Malformed("block_quantize block_size must be greater than zero")
353
+ if self.levels == 0 or self.levels > MAX_BLOCK_LEVELS:
354
+ raise Malformed(
355
+ "block_quantize levels %d is outside 1..=%d"
356
+ % (self.levels, MAX_BLOCK_LEVELS)
357
+ )
358
+ if element_count % self.block_size != 0:
359
+ raise Malformed(
360
+ "block_quantize block_size %d does not divide element_count %d"
361
+ % (self.block_size, element_count)
362
+ )
363
+
364
+
365
+ def parse_semantic_policy(value, dtype: str, element_count: int) -> SemanticPolicy:
366
+ map_value = value if isinstance(value, dict) else None
367
+ if map_value is None:
368
+ raise Malformed("semantic policy must be a WCB map")
369
+ scheme = map_value.get("scheme")
370
+ if not isinstance(scheme, str):
371
+ raise Malformed("semantic policy `scheme` must be text")
372
+ if scheme == "none":
373
+ _exact_keys(map_value, ("scheme",), "semantic policy")
374
+ return SemanticPolicy("none")
375
+ if scheme == "mantissa_truncate":
376
+ _exact_keys(map_value, ("bits", "scheme"), "semantic policy")
377
+ policy = SemanticPolicy("mantissa_truncate", bits=_u64(map_value, "bits"))
378
+ policy.validate(dtype, element_count)
379
+ return policy
380
+ if scheme == "block_quantize":
381
+ _exact_keys(map_value, ("block_size", "levels", "scheme"), "semantic policy")
382
+ policy = SemanticPolicy(
383
+ "block_quantize",
384
+ block_size=_u64(map_value, "block_size"),
385
+ levels=_u64(map_value, "levels"),
386
+ )
387
+ policy.validate(dtype, element_count)
388
+ return policy
389
+ raise Malformed(
390
+ "unknown semantic policy scheme `%s` (v1 allows none, mantissa_truncate, "
391
+ "block_quantize)" % scheme
392
+ )
393
+
394
+
395
+ def _signed_word(dtype: str, word: int) -> int:
396
+ bits = WEIGHT_WORD_BITS[dtype]
397
+ sign = 1 << (bits - 1)
398
+ modulus = 1 << bits
399
+ if word >= sign:
400
+ return word - modulus
401
+ return word
402
+
403
+
404
+ def _raw_word(dtype: str, signed: int) -> int:
405
+ return signed & ((1 << WEIGHT_WORD_BITS[dtype]) - 1)
406
+
407
+
408
+ def _exponent_all_ones(dtype: str, word: int) -> bool:
409
+ if dtype in ("bf16", "fp16"):
410
+ return (word & 0x7F80) == 0x7F80
411
+ if dtype == "fp32":
412
+ return (word & 0x7F800000) == 0x7F800000
413
+ return False
414
+
415
+
416
+ def _word_at(dtype: str, data: bytes, index: int) -> int:
417
+ if dtype in ("bf16", "fp16"):
418
+ offset = index * 2
419
+ return int.from_bytes(data[offset : offset + 2], "little")
420
+ if dtype == "fp32":
421
+ offset = index * 4
422
+ return int.from_bytes(data[offset : offset + 4], "little")
423
+ return data[index]
424
+
425
+
426
+ def _push_word(dtype: str, word: int, out: bytearray) -> None:
427
+ if dtype in ("bf16", "fp16"):
428
+ out.extend(int(word & 0xFFFF).to_bytes(2, "little"))
429
+ elif dtype == "fp32":
430
+ out.extend(int(word & 0xFFFFFFFF).to_bytes(4, "little"))
431
+ else:
432
+ out.append(word & 0xFF)
433
+
434
+
435
+ def _transform_bytes(
436
+ dtype: str, data: bytes, element_count: int, policy: SemanticPolicy
437
+ ) -> bytes:
438
+ if policy.scheme == "none":
439
+ return bytes(data)
440
+ if policy.scheme == "mantissa_truncate":
441
+ retain = ~((1 << policy.bits) - 1) & ((1 << WEIGHT_WORD_BITS[dtype]) - 1)
442
+ out = bytearray()
443
+ for index in range(element_count):
444
+ word = _word_at(dtype, data, index)
445
+ transformed = (
446
+ word if _exponent_all_ones(dtype, word) else word & retain
447
+ )
448
+ _push_word(dtype, transformed, out)
449
+ return bytes(out)
450
+ block = policy.block_size
451
+ levels = policy.levels
452
+ out = bytearray()
453
+ start = 0
454
+ while start < element_count:
455
+ end = min(start + block, element_count)
456
+ amax = 0
457
+ for index in range(start, end):
458
+ magnitude = abs(_signed_word(dtype, _word_at(dtype, data, index)))
459
+ if magnitude > amax:
460
+ amax = magnitude
461
+ if amax == 0 or levels == 1:
462
+ for _ in range(start, end):
463
+ _push_word(dtype, 0, out)
464
+ else:
465
+ span = 2 * amax
466
+ for index in range(start, end):
467
+ value = _signed_word(dtype, _word_at(dtype, data, index))
468
+ numerator = (value + amax) * (levels - 1) + amax
469
+ k = numerator // span
470
+ if k >= levels:
471
+ k = levels - 1
472
+ level = (span * k + (levels - 1) // 2) // (levels - 1)
473
+ _push_word(dtype, _raw_word(dtype, level - amax), out)
474
+ start = end
475
+ return bytes(out)
476
+
477
+
478
+ def semantic_root(
479
+ dtype: str, shape: "list[int]", data: bytes, policy: SemanticPolicy, limits: Limits
480
+ ) -> bytes:
481
+ count = _validate_shape(shape, limits, "semantic tensor")
482
+ expected = _checked_mul(
483
+ count, WEIGHT_STORAGE_BYTES[dtype], "semantic tensor byte length"
484
+ )
485
+ _check_len(expected, limits.max_decoded_bytes, "semantic tensor bytes")
486
+ if expected != len(data):
487
+ raise Malformed(
488
+ "semantic tensor is %d bytes but shape %r x %s requires %d"
489
+ % (len(data), shape, dtype, expected)
490
+ )
491
+ policy.validate(dtype, count)
492
+ if dtype == "q4_0":
493
+ for byte in data:
494
+ if byte > 15:
495
+ raise Corrupt("q4_0 semantic tensor word is outside 0..=15")
496
+ if policy.scheme == "none":
497
+ quantized_values_hash = values_hash(bytes(data))
498
+ else:
499
+ quantized_values_hash = hash_domain(
500
+ DOMAIN_TENSOR_SEMANTIC_VALUES,
501
+ _transform_bytes(dtype, data, count, policy),
502
+ )
503
+ record = {
504
+ "schema": SEMANTIC_SCHEMA,
505
+ "dtype": dtype,
506
+ "shape": _u64_vec(shape),
507
+ "element_count": wcb.U64(count),
508
+ "policy": policy.to_wcb(),
509
+ "quantized_values_hash": quantized_values_hash,
510
+ }
511
+ return hash_domain(DOMAIN_TENSOR_SEMANTIC, wcb.encode(record, limits))
512
+
513
+
514
+ def semantic_record_from_canonical(data: bytes, limits: Limits) -> dict:
515
+ """Parse and validate a canonical ``weft.tensor.semantic.v1`` record."""
516
+ value = wcb.decode(data, limits)
517
+ map_value = _map(value, "semantic root record")
518
+ _exact_keys(
519
+ map_value,
520
+ ("dtype", "element_count", "policy", "quantized_values_hash", "schema", "shape"),
521
+ "semantic root record",
522
+ )
523
+ schema = _text(map_value, "schema")
524
+ if schema != SEMANTIC_SCHEMA:
525
+ raise Malformed("unsupported semantic root schema `%s`" % schema)
526
+ dtype = _weight_dtype(_text(map_value, "dtype"))
527
+ shape = _shape(map_value)
528
+ element_count = _u64(map_value, "element_count")
529
+ policy = parse_semantic_policy(map_value.get("policy"), dtype, element_count)
530
+ quantized_values_hash = _hash_field(map_value, "quantized_values_hash")
531
+ count = _validate_shape(shape, limits, "semantic")
532
+ if element_count != count:
533
+ raise Malformed(
534
+ "semantic element_count %d does not match shape product %d"
535
+ % (element_count, count)
536
+ )
537
+ _checked_mul(count, WEIGHT_STORAGE_BYTES[dtype], "semantic tensor byte length")
538
+ policy.validate(dtype, count)
539
+ record = {
540
+ "schema": schema,
541
+ "dtype": dtype,
542
+ "shape": _u64_vec(shape),
543
+ "element_count": wcb.U64(element_count),
544
+ "policy": policy.to_wcb(),
545
+ "quantized_values_hash": quantized_values_hash,
546
+ }
547
+ if wcb.encode(record, limits) != data:
548
+ raise NonCanonical("semantic root bytes are not in canonical form")
549
+ return record
550
+
551
+
552
+ # ---------------------------------------------------------------------------
553
+ # Mode 9 -- FLOAT_TENSOR
554
+ # ---------------------------------------------------------------------------
555
+
556
+ _FLOAT_KEYS = (
557
+ "schema",
558
+ "dtype",
559
+ "shape",
560
+ "element_count",
561
+ "sign_bits",
562
+ "exponent_bits",
563
+ "mantissa_bits",
564
+ "exponent_coder",
565
+ "exponent_plane_len",
566
+ "mantissa_plane_len",
567
+ "payload_hash",
568
+ )
569
+
570
+
571
+ @dataclass(frozen=True)
572
+ class FloatTensorMeta:
573
+ dtype: str
574
+ shape: "list[int]"
575
+ element_count: int
576
+ sign_bits: int
577
+ exponent_bits: int
578
+ mantissa_bits: int
579
+ exponent_plane_len: int
580
+ mantissa_plane_len: int
581
+
582
+
583
+ def _float_dtype(name: str) -> str:
584
+ if name not in FLOAT_FIELD_BITS:
585
+ raise Malformed(
586
+ "unknown float tensor dtype `%s` (v1 allows bf16, fp16, fp32, "
587
+ "fp8-e4m3, fp8-e5m2)" % name
588
+ )
589
+ return name
590
+
591
+
592
+ def _byte_field(map_value: dict, key: str) -> int:
593
+ value = _u64(map_value, key)
594
+ if value > 0xFF:
595
+ raise Malformed(
596
+ "float tensor `%s` must fit an unsigned byte, got %d" % (key, value)
597
+ )
598
+ return value
599
+
600
+
601
+ def parse_float_descriptor(
602
+ frame, limits: Limits
603
+ ) -> "tuple[FloatTensorMeta, bytes, bytes, bytes, bytes]":
604
+ descriptor_bytes, stored = _split_prefix(frame.body)
605
+ map_value = _map(_decode_descriptor(descriptor_bytes, limits), "float tensor")
606
+ _exact_keys(map_value, _FLOAT_KEYS, "float tensor")
607
+ schema = _text(map_value, "schema")
608
+ if schema != FLOAT_TENSOR_SCHEMA:
609
+ raise Malformed("unsupported float tensor schema `%s`" % schema)
610
+ dtype = _float_dtype(_text(map_value, "dtype"))
611
+ coder = _text(map_value, "exponent_coder")
612
+ if coder != PLANE_CODEC:
613
+ raise Malformed(
614
+ "unsupported float tensor exponent coder `%s` (v1 allows `%s`)"
615
+ % (coder, PLANE_CODEC)
616
+ )
617
+ sign_bits = _byte_field(map_value, "sign_bits")
618
+ exponent_bits = _byte_field(map_value, "exponent_bits")
619
+ mantissa_bits = _byte_field(map_value, "mantissa_bits")
620
+ if (sign_bits, exponent_bits, mantissa_bits) != FLOAT_FIELD_BITS[dtype]:
621
+ raise Malformed(
622
+ "float tensor sub-field widths %d, %d, %d disagree with dtype `%s`"
623
+ % (sign_bits, exponent_bits, mantissa_bits, dtype)
624
+ )
625
+ element_bits = sign_bits + exponent_bits + mantissa_bits
626
+ if element_bits % 8 != 0:
627
+ raise Malformed("float tensor element width is not byte-aligned")
628
+ element_size = element_bits // 8
629
+ shape = _shape(map_value)
630
+ declared_count = _u64(map_value, "element_count")
631
+ expected_count = _validate_shape(shape, limits, "float tensor")
632
+ if declared_count != expected_count:
633
+ raise Malformed(
634
+ "float tensor element_count %d does not match shape product %d"
635
+ % (declared_count, expected_count)
636
+ )
637
+ decoded_bytes = _checked_mul(
638
+ declared_count, element_size, "float tensor decoded byte count"
639
+ )
640
+ _check_len(decoded_bytes, limits.max_decoded_bytes, "float tensor decoded bytes")
641
+ if decoded_bytes != frame.header.decoded_len:
642
+ raise Corrupt(
643
+ "float tensor decoded_len is %d but the descriptor requires %d"
644
+ % (frame.header.decoded_len, decoded_bytes)
645
+ )
646
+ mantissa_bits_total = _checked_mul(
647
+ declared_count,
648
+ sign_bits + mantissa_bits,
649
+ "float tensor mantissa bit count",
650
+ )
651
+ expected_mantissa = _checked_add(
652
+ mantissa_bits_total, 7, "float tensor mantissa length"
653
+ ) // 8
654
+ declared_mantissa = _u64(map_value, "mantissa_plane_len")
655
+ if declared_mantissa != expected_mantissa:
656
+ raise Malformed(
657
+ "float tensor mantissa_plane_len %d does not match element_count x "
658
+ "(sign_bits + mantissa_bits) / 8 %d"
659
+ % (declared_mantissa, expected_mantissa)
660
+ )
661
+ declared_exponent = _u64(map_value, "exponent_plane_len")
662
+ stored_len = _checked_add(
663
+ declared_exponent, declared_mantissa, "float tensor plane lengths"
664
+ )
665
+ if stored_len != len(stored):
666
+ raise Malformed(
667
+ "float tensor stored payload is %d bytes but the descriptor declares %d"
668
+ % (len(stored), stored_len)
669
+ )
670
+ exponent_compressed = stored[:declared_exponent]
671
+ mantissa_plane = stored[declared_exponent:]
672
+ meta = FloatTensorMeta(
673
+ dtype=dtype,
674
+ shape=shape,
675
+ element_count=declared_count,
676
+ sign_bits=sign_bits,
677
+ exponent_bits=exponent_bits,
678
+ mantissa_bits=mantissa_bits,
679
+ exponent_plane_len=declared_exponent,
680
+ mantissa_plane_len=declared_mantissa,
681
+ )
682
+ return meta, stored, exponent_compressed, mantissa_plane, _bytes_of(
683
+ map_value, "payload_hash"
684
+ )
685
+
686
+
687
+ def _join_float_planes(
688
+ exponent_plane: bytes,
689
+ mantissa_plane: bytes,
690
+ dtype: str,
691
+ element_count: int,
692
+ limits: Limits,
693
+ ) -> bytes:
694
+ if len(exponent_plane) != element_count:
695
+ raise Corrupt(
696
+ "float tensor exponent plane has %d bytes but element_count is %d"
697
+ % (len(exponent_plane), element_count)
698
+ )
699
+ sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[dtype]
700
+ element_size = (sign_bits + exponent_bits + mantissa_bits) // 8
701
+ decoded_bytes = _checked_mul(
702
+ element_count, element_size, "float tensor decoded byte count"
703
+ )
704
+ _check_len(decoded_bytes, limits.max_decoded_bytes, "float tensor decoded bytes")
705
+ expected_mantissa_bits = _checked_mul(
706
+ element_count,
707
+ sign_bits + mantissa_bits,
708
+ "float tensor mantissa bit count",
709
+ )
710
+ expected_mantissa_bytes = _checked_add(
711
+ expected_mantissa_bits, 7, "float tensor mantissa length"
712
+ ) // 8
713
+ if len(mantissa_plane) != expected_mantissa_bytes:
714
+ raise Corrupt(
715
+ "float tensor mantissa plane has %d bytes but descriptor arithmetic "
716
+ "requires %d" % (len(mantissa_plane), expected_mantissa_bytes)
717
+ )
718
+
719
+ width = sign_bits + mantissa_bits
720
+ bit = 0
721
+ total_bits = len(mantissa_plane) * 8
722
+ out = bytearray(decoded_bytes)
723
+ mantissa_mask = (1 << mantissa_bits) - 1
724
+ for index, exponent_byte in enumerate(exponent_plane):
725
+ if exponent_byte >> exponent_bits:
726
+ raise Corrupt(
727
+ "float tensor exponent byte exceeds the dtype exponent width"
728
+ )
729
+ value = 0
730
+ for _ in range(width):
731
+ if bit >= total_bits:
732
+ # The plane is exactly long enough by construction; this is an
733
+ # internal invariant, reported as corruption rather than panic.
734
+ raise Corrupt("float tensor mantissa plane is truncated")
735
+ byte = mantissa_plane[bit >> 3]
736
+ value = (value << 1) | ((byte >> (7 - (bit & 7))) & 1)
737
+ bit += 1
738
+ sign = value >> mantissa_bits
739
+ mantissa = value & mantissa_mask
740
+ word = (
741
+ (sign << (exponent_bits + mantissa_bits))
742
+ | (exponent_byte << mantissa_bits)
743
+ | mantissa
744
+ )
745
+ offset = index * element_size
746
+ out[offset : offset + element_size] = word.to_bytes(element_size, "little")
747
+
748
+ padding = len(mantissa_plane) * 8 - expected_mantissa_bits
749
+ if padding > 0:
750
+ last = mantissa_plane[-1]
751
+ if last & ((1 << padding) - 1):
752
+ raise Corrupt("float tensor mantissa padding bits must be zero")
753
+ return bytes(out)
754
+
755
+
756
+ def decode_float_tensor(frame, limits: Limits) -> bytes:
757
+ meta, stored, exponent_compressed, mantissa_plane, declared_hash = (
758
+ parse_float_descriptor(frame, limits)
759
+ )
760
+ exponent_plane = huffman.decode(exponent_compressed, meta.element_count)
761
+ out = _join_float_planes(
762
+ exponent_plane,
763
+ mantissa_plane,
764
+ meta.dtype,
765
+ meta.element_count,
766
+ limits,
767
+ )
768
+ if len(declared_hash) != HASH_LEN:
769
+ raise Malformed(
770
+ "float tensor payload_hash must be %d bytes, got %d"
771
+ % (HASH_LEN, len(declared_hash))
772
+ )
773
+ if declared_hash != hash_domain("weft/float-tensor/payload/v1", stored):
774
+ raise Corrupt("float tensor payload_hash mismatch")
775
+ return out
776
+
777
+
778
+ # ---------------------------------------------------------------------------
779
+ # Mode 10 -- TENSOR_SHARDED
780
+ # ---------------------------------------------------------------------------
781
+
782
+ _SHARDED_KEYS = (
783
+ "schema",
784
+ "dtype",
785
+ "shape",
786
+ "element_count",
787
+ "shard_axis",
788
+ "shard_count",
789
+ "shard_index",
790
+ "block_size",
791
+ "scales",
792
+ "zero_points",
793
+ "layout_root",
794
+ "payload_len",
795
+ "payload_hash",
796
+ )
797
+
798
+
799
+ @dataclass
800
+ class ShardedMeta:
801
+ dtype: str
802
+ shape: "list[int]"
803
+ element_count: int
804
+ shard_axis: int
805
+ shard_count: int
806
+ shard_index: int
807
+ block_size: int
808
+ scales: "list[int]"
809
+ zero_points: "list[int]"
810
+ layout_root: bytes
811
+ payload_len: int
812
+
813
+
814
+ @dataclass
815
+ class DecodedShard:
816
+ meta: ShardedMeta
817
+ words: "list[int]"
818
+ planes: bytes
819
+
820
+
821
+ def _sharded_dtype(name: str) -> str:
822
+ if name not in ("q8_0", "q4_0"):
823
+ raise Malformed(
824
+ "unknown TENSOR_SHARDED dtype `%s` (v1 allows q8_0, q4_0)" % name
825
+ )
826
+ return name
827
+
828
+
829
+ def _u32_array(map_value: dict, key: str) -> "list[int]":
830
+ values = _array_u64(map_value, key)
831
+ for value in values:
832
+ if value > U32_MAX:
833
+ raise Malformed(
834
+ "TENSOR_SHARDED scale bits exceed the 32-bit f32 pattern"
835
+ )
836
+ return values
837
+
838
+
839
+ def _i64_array(map_value: dict, key: str) -> "list[int]":
840
+ value = map_value.get(key)
841
+ if not isinstance(value, list):
842
+ raise Malformed("metadata `%s` must be an array" % key)
843
+ out = []
844
+ for item in value:
845
+ if not isinstance(item, wcb.I64):
846
+ raise Malformed("metadata `%s` elements must be i64" % key)
847
+ out.append(item.value)
848
+ return out
849
+
850
+
851
+ def parse_sharded_descriptor(descriptor_bytes: bytes, limits: Limits):
852
+ map_value = _map(_decode_descriptor(descriptor_bytes, limits), "TENSOR_SHARDED")
853
+ _exact_keys(map_value, _SHARDED_KEYS, "TENSOR_SHARDED")
854
+ schema = _text(map_value, "schema")
855
+ if schema != TENSOR_SHARDED_SCHEMA:
856
+ raise Malformed("unsupported TENSOR_SHARDED schema `%s`" % schema)
857
+ dtype = _sharded_dtype(_text(map_value, "dtype"))
858
+ shape = _shape(map_value)
859
+ element_count = _u64(map_value, "element_count")
860
+ shard_axis = _u32(map_value, "shard_axis")
861
+ shard_count = _u32(map_value, "shard_count")
862
+ shard_index = _u32(map_value, "shard_index")
863
+ block_size = _u64(map_value, "block_size")
864
+ scales = _u32_array(map_value, "scales")
865
+ zero_points = _i64_array(map_value, "zero_points")
866
+ layout_root = _hash_field(map_value, "layout_root")
867
+ payload_len = _u64(map_value, "payload_len")
868
+ payload_hash = _hash_field(map_value, "payload_hash")
869
+ meta = ShardedMeta(
870
+ dtype=dtype,
871
+ shape=shape,
872
+ element_count=element_count,
873
+ shard_axis=shard_axis,
874
+ shard_count=shard_count,
875
+ shard_index=shard_index,
876
+ block_size=block_size,
877
+ scales=scales,
878
+ zero_points=zero_points,
879
+ layout_root=layout_root,
880
+ payload_len=payload_len,
881
+ )
882
+ return meta, payload_hash
883
+
884
+
885
+ def _validate_sharded_meta(
886
+ meta: ShardedMeta, limits: Limits
887
+ ) -> "tuple[int, int, int]":
888
+ element_count = _validate_shape(meta.shape, limits, "TENSOR_SHARDED")
889
+ if meta.element_count != element_count:
890
+ raise Malformed(
891
+ "TENSOR_SHARDED element_count %d does not match shape product %d"
892
+ % (meta.element_count, element_count)
893
+ )
894
+ if meta.shard_count == 0:
895
+ raise Malformed("TENSOR_SHARDED shard_count must be at least 1")
896
+ if meta.shard_axis >= len(meta.shape):
897
+ raise Malformed(
898
+ "TENSOR_SHARDED shard_axis %d is not below rank %d"
899
+ % (meta.shard_axis, len(meta.shape))
900
+ )
901
+ if meta.shard_index >= meta.shard_count:
902
+ raise Malformed(
903
+ "TENSOR_SHARDED shard_index %d is not below shard_count %d"
904
+ % (meta.shard_index, meta.shard_count)
905
+ )
906
+ axis_dim = meta.shape[meta.shard_axis]
907
+ if axis_dim % meta.shard_count != 0:
908
+ raise Malformed(
909
+ "TENSOR_SHARDED shape[%d] = %d is not divisible by shard_count %d"
910
+ % (meta.shard_axis, axis_dim, meta.shard_count)
911
+ )
912
+ if element_count % meta.shard_count != 0:
913
+ raise Malformed(
914
+ "TENSOR_SHARDED element_count %d is not divisible by shard_count %d"
915
+ % (element_count, meta.shard_count)
916
+ )
917
+ shard_elements = element_count // meta.shard_count
918
+ if meta.block_size not in BLOCK_SIZES:
919
+ raise Malformed(
920
+ "TENSOR_SHARDED block_size %d is not one of 16, 32, 64, 128, 256"
921
+ % meta.block_size
922
+ )
923
+ if element_count % meta.block_size != 0:
924
+ raise Malformed(
925
+ "TENSOR_SHARDED element_count %d is not divisible by block_size %d"
926
+ % (element_count, meta.block_size)
927
+ )
928
+ block_count = element_count // meta.block_size
929
+ if block_count > limits.max_structured_items:
930
+ raise LimitExceeded(
931
+ "TENSOR_SHARDED block count %d exceeds WCB item limit %d"
932
+ % (block_count, limits.max_structured_items)
933
+ )
934
+ if len(meta.scales) != block_count or len(meta.zero_points) != block_count:
935
+ raise Malformed(
936
+ "TENSOR_SHARDED descriptor has %d scales and %d zero points but "
937
+ "requires %d"
938
+ % (len(meta.scales), len(meta.zero_points), block_count)
939
+ )
940
+ for bits in meta.scales:
941
+ scale = struct.unpack("<f", struct.pack("<I", bits))[0]
942
+ if not math.isfinite(scale):
943
+ raise Malformed("TENSOR_SHARDED scale bits must be a finite f32")
944
+ if scale <= 0.0:
945
+ raise Malformed("TENSOR_SHARDED scale must be greater than zero")
946
+ if meta.dtype == "q4_0" and shard_elements % 2 != 0:
947
+ raise Malformed(
948
+ "q4_0 shard element count must be even to pack nibble pairs"
949
+ )
950
+ quant_len = shard_elements if meta.dtype == "q8_0" else shard_elements // 2
951
+ planes_len = _checked_add(
952
+ quant_len, _checked_mul(block_count, 4, "TENSOR_SHARDED scale plane"),
953
+ "TENSOR_SHARDED plane length",
954
+ )
955
+ _check_len(
956
+ planes_len, limits.max_decoded_bytes, "TENSOR_SHARDED reconstructed planes"
957
+ )
958
+ return element_count, shard_elements, block_count
959
+
960
+
961
+ def decode_tensor_sharded(frame, limits: Limits) -> DecodedShard:
962
+ descriptor_bytes, stored = _split_prefix(frame.body)
963
+ meta, declared_payload_hash = parse_sharded_descriptor(descriptor_bytes, limits)
964
+ _element_count, shard_elements, block_count = _validate_sharded_meta(meta, limits)
965
+ if meta.payload_len != len(stored):
966
+ raise Malformed(
967
+ "TENSOR_SHARDED payload is %d bytes but payload_len declares %d"
968
+ % (len(stored), meta.payload_len)
969
+ )
970
+ quant_len = shard_elements if meta.dtype == "q8_0" else shard_elements // 2
971
+ planes_len = _checked_add(
972
+ quant_len, _checked_mul(block_count, 4, "TENSOR_SHARDED scale plane"),
973
+ "TENSOR_SHARDED plane length",
974
+ )
975
+ if frame.header.decoded_len != planes_len:
976
+ raise Corrupt(
977
+ "TENSOR_SHARDED decoded_len is %d but the descriptor requires %d"
978
+ % (frame.header.decoded_len, planes_len)
979
+ )
980
+ if declared_payload_hash != hash_domain(DOMAIN_TENSOR_SHARDED_PAYLOAD, stored):
981
+ raise Corrupt("TENSOR_SHARDED payload_hash mismatch")
982
+ _check_len(
983
+ planes_len, limits.max_decoded_bytes, "TENSOR_SHARDED reconstructed planes"
984
+ )
985
+ planes = huffman.decode(stored, planes_len)
986
+ if quant_len > len(planes):
987
+ raise Corrupt(
988
+ "TENSOR_SHARDED reconstructed plane length is inconsistent"
989
+ )
990
+ quant_plane = planes[:quant_len]
991
+ scale_bytes = planes[quant_len:]
992
+ if len(scale_bytes) != block_count * 4:
993
+ raise Corrupt("TENSOR_SHARDED scale plane length is inconsistent")
994
+ for index in range(block_count):
995
+ chunk = scale_bytes[index * 4 : index * 4 + 4]
996
+ actual = int.from_bytes(chunk, "little")
997
+ if actual != meta.scales[index]:
998
+ raise Corrupt(
999
+ "TENSOR_SHARDED scale plane block %d does not match the descriptor"
1000
+ % index
1001
+ )
1002
+ if frame.header.content_hash != hash_domain("weft/chunk/v1", planes):
1003
+ raise Corrupt(
1004
+ "TENSOR_SHARDED content_hash mismatch against the reconstructed planes"
1005
+ )
1006
+ if meta.dtype == "q8_0":
1007
+ words = [byte - 256 if byte >= 128 else byte for byte in quant_plane]
1008
+ else:
1009
+ words = []
1010
+ for byte in quant_plane:
1011
+ low = byte & 0x0F
1012
+ high = byte >> 4
1013
+ words.append(low - 16 if low >= 8 else low)
1014
+ words.append(high - 16 if high >= 8 else high)
1015
+ words = words[:shard_elements]
1016
+ if len(words) != shard_elements:
1017
+ raise Corrupt(
1018
+ "TENSOR_SHARDED reconstructed word count is inconsistent"
1019
+ )
1020
+ return DecodedShard(meta=meta, words=words, planes=planes)
1021
+
1022
+
1023
+ def _same_sharded_tensor(left: ShardedMeta, right: ShardedMeta) -> bool:
1024
+ return (
1025
+ left.dtype == right.dtype
1026
+ and left.shape == right.shape
1027
+ and left.element_count == right.element_count
1028
+ and left.shard_axis == right.shard_axis
1029
+ and left.shard_count == right.shard_count
1030
+ and left.block_size == right.block_size
1031
+ and left.scales == right.scales
1032
+ and left.zero_points == right.zero_points
1033
+ and left.layout_root == right.layout_root
1034
+ )
1035
+
1036
+
1037
+ def _assemble_shard_words(
1038
+ shape: "list[int]", axis: int, shards: "list[list[int]]"
1039
+ ) -> "list[int]":
1040
+ count = len(shards)
1041
+ shard_shape = list(shape)
1042
+ shard_shape[axis] = shard_shape[axis] // count
1043
+ total = 1
1044
+ for dim in shape:
1045
+ total *= dim
1046
+ out = [0] * total
1047
+ plane = shape[axis] // count
1048
+ strides = [0] * len(shape)
1049
+ stride = 1
1050
+ for index in range(len(shape) - 1, -1, -1):
1051
+ strides[index] = stride
1052
+ stride *= shape[index]
1053
+ for shard_index, payload in enumerate(shards):
1054
+ # Iterate the shard-local row-major coordinates.
1055
+ def walk(dimension: int, local_coord: "list[int]", local_flat: int) -> int:
1056
+ if dimension == len(shard_shape):
1057
+ full = 0
1058
+ for axis_index, coord in enumerate(local_coord):
1059
+ full_coord = coord
1060
+ if axis_index == axis:
1061
+ full_coord = shard_index * plane + coord
1062
+ full += full_coord * strides[axis_index]
1063
+ out[full] = payload[local_flat]
1064
+ return local_flat + 1
1065
+ extent = shard_shape[dimension]
1066
+ for position in range(extent):
1067
+ local_coord[dimension] = position
1068
+ local_flat = walk(dimension + 1, local_coord, local_flat)
1069
+ return local_flat
1070
+
1071
+ walk(0, [0] * len(shard_shape), 0)
1072
+ return out
1073
+
1074
+
1075
+ def shard_set_root(
1076
+ shards: "list[DecodedShard]", limits: Limits
1077
+ ) -> "tuple[bytes, list[int]]":
1078
+ if not shards:
1079
+ raise Malformed("TENSOR_SHARDED shard set is empty")
1080
+ first = shards[0].meta
1081
+ _validate_sharded_meta(first, limits)
1082
+ count = first.shard_count
1083
+ if len(shards) != count:
1084
+ raise Malformed(
1085
+ "TENSOR_SHARDED shard set has %d entries but the descriptor declares %d"
1086
+ % (len(shards), count)
1087
+ )
1088
+ ordered: "list[DecodedShard | None]" = [None] * count
1089
+ for shard in shards:
1090
+ if not _same_sharded_tensor(first, shard.meta):
1091
+ raise NonCanonical(
1092
+ "TENSOR_SHARDED shard descriptors disagree on tensor identity"
1093
+ )
1094
+ index = shard.meta.shard_index
1095
+ if index >= count:
1096
+ raise NonCanonical(
1097
+ "TENSOR_SHARDED shard index %d is not below shard count %d"
1098
+ % (index, count)
1099
+ )
1100
+ if ordered[index] is not None:
1101
+ raise NonCanonical(
1102
+ "duplicate TENSOR_SHARDED shard index %d" % index
1103
+ )
1104
+ ordered[index] = shard
1105
+ for index, entry in enumerate(ordered):
1106
+ if entry is None:
1107
+ raise NonCanonical("missing TENSOR_SHARDED shard index %d" % index)
1108
+ words = _assemble_shard_words(
1109
+ first.shape, first.shard_axis, [entry.words for entry in ordered]
1110
+ )
1111
+ root = sharded_logical_root(
1112
+ first.dtype,
1113
+ first.shape,
1114
+ words,
1115
+ first.scales,
1116
+ first.zero_points,
1117
+ limits,
1118
+ )
1119
+ if root != first.layout_root:
1120
+ raise Corrupt(
1121
+ "TENSOR_SHARDED shard set recomputes a different logical root"
1122
+ )
1123
+ return root, words
1124
+
1125
+
1126
+ # ---------------------------------------------------------------------------
1127
+ # Modes 11/14/16 -- delta modes
1128
+ # ---------------------------------------------------------------------------
1129
+
1130
+
1131
+ def _weight_dtype(name: str) -> str:
1132
+ if name not in WEIGHT_WORD_BITS:
1133
+ raise Malformed(
1134
+ "unknown WEIGHT_DELTA dtype `%s` (v1 allows bf16, fp16, fp32, "
1135
+ "q8_0, q4_0)" % name
1136
+ )
1137
+ return name
1138
+
1139
+
1140
+ @dataclass
1141
+ class WeightDeltaMeta:
1142
+ dtype: str
1143
+ shape: "list[int]"
1144
+ element_count: int
1145
+ word_bits: int
1146
+ base: bytes
1147
+ base_dtype: str
1148
+ base_shape: "list[int]"
1149
+ op: str
1150
+ planes: int
1151
+ plane_codec: str
1152
+ delta_root: bytes
1153
+ payload_len: int
1154
+
1155
+
1156
+ _WEIGHT_KEYS = (
1157
+ "schema",
1158
+ "dtype",
1159
+ "shape",
1160
+ "element_count",
1161
+ "word_bits",
1162
+ "base",
1163
+ "base_dtype",
1164
+ "base_shape",
1165
+ "op",
1166
+ "planes",
1167
+ "plane_codec",
1168
+ "delta_root",
1169
+ "payload_len",
1170
+ )
1171
+
1172
+
1173
+ def parse_weight_descriptor(descriptor_bytes: bytes, limits: Limits) -> WeightDeltaMeta:
1174
+ map_value = _map(_decode_descriptor(descriptor_bytes, limits), "WEIGHT_DELTA")
1175
+ _exact_keys(map_value, _WEIGHT_KEYS, "WEIGHT_DELTA")
1176
+ schema = _text(map_value, "schema")
1177
+ if schema != WEIGHT_DELTA_SCHEMA:
1178
+ raise Malformed("unsupported WEIGHT_DELTA schema `%s`" % schema)
1179
+ dtype = _weight_dtype(_text(map_value, "dtype"))
1180
+ shape = _shape(map_value)
1181
+ element_count = _u64(map_value, "element_count")
1182
+ word_bits = _u64(map_value, "word_bits")
1183
+ base = _hash_field(map_value, "base")
1184
+ base_dtype = _weight_dtype(_text(map_value, "base_dtype"))
1185
+ base_shape = _array_u64(map_value, "base_shape")
1186
+ op = _text(map_value, "op")
1187
+ planes = _u64(map_value, "planes")
1188
+ plane_codec = _text(map_value, "plane_codec")
1189
+ delta_root = _hash_field(map_value, "delta_root")
1190
+ payload_len = _u64(map_value, "payload_len")
1191
+ meta = WeightDeltaMeta(
1192
+ dtype=dtype,
1193
+ shape=shape,
1194
+ element_count=element_count,
1195
+ word_bits=word_bits,
1196
+ base=base,
1197
+ base_dtype=base_dtype,
1198
+ base_shape=base_shape,
1199
+ op=op,
1200
+ planes=planes,
1201
+ plane_codec=plane_codec,
1202
+ delta_root=delta_root,
1203
+ payload_len=payload_len,
1204
+ )
1205
+ validate_weight_meta(meta, limits)
1206
+ return meta
1207
+
1208
+
1209
+ def validate_weight_meta(meta: WeightDeltaMeta, limits: Limits) -> int:
1210
+ if meta.op != DELTA_OP:
1211
+ raise Malformed(
1212
+ "WEIGHT_DELTA op `%s` is not implemented (only %s)"
1213
+ % (meta.op, DELTA_OP)
1214
+ )
1215
+ if meta.plane_codec != PLANE_CODEC:
1216
+ raise Malformed(
1217
+ "WEIGHT_DELTA plane_codec `%s` is not implemented (only %s)"
1218
+ % (meta.plane_codec, PLANE_CODEC)
1219
+ )
1220
+ if WEIGHT_WORD_BITS[meta.dtype] != meta.word_bits:
1221
+ raise Malformed(
1222
+ "WEIGHT_DELTA word_bits %d disagrees with dtype %s (%d bits)"
1223
+ % (meta.word_bits, meta.dtype, WEIGHT_WORD_BITS[meta.dtype])
1224
+ )
1225
+ if WEIGHT_PLANES[meta.dtype] != meta.planes:
1226
+ raise Malformed(
1227
+ "WEIGHT_DELTA planes %d disagrees with dtype %s (%d planes)"
1228
+ % (meta.planes, meta.dtype, WEIGHT_PLANES[meta.dtype])
1229
+ )
1230
+ if meta.base_dtype != meta.dtype:
1231
+ raise Malformed(
1232
+ "WEIGHT_DELTA base_dtype `%s` must equal dtype `%s`"
1233
+ % (meta.base_dtype, meta.dtype)
1234
+ )
1235
+ if meta.base_shape != meta.shape:
1236
+ raise Malformed("WEIGHT_DELTA base_shape must equal shape")
1237
+ count = _validate_shape(meta.shape, limits, "WEIGHT_DELTA")
1238
+ if meta.element_count != count:
1239
+ raise Malformed(
1240
+ "WEIGHT_DELTA element_count %d does not match shape product %d"
1241
+ % (meta.element_count, count)
1242
+ )
1243
+ _checked_mul(
1244
+ count, WEIGHT_STORAGE_BYTES[meta.dtype],
1245
+ "WEIGHT_DELTA element byte length",
1246
+ )
1247
+ return count
1248
+
1249
+
1250
+ def weight_metadata(frame, limits: Limits) -> WeightDeltaMeta:
1251
+ descriptor_bytes, _payload = _split_prefix(frame.body)
1252
+ return parse_weight_descriptor(descriptor_bytes, limits)
1253
+
1254
+
1255
+ def _decode_planes(dtype: str, payload: bytes, element_count: int, what: str):
1256
+ plane_count = WEIGHT_PLANES[dtype]
1257
+ planes = []
1258
+ offset = 0
1259
+ for _ in range(plane_count):
1260
+ if offset >= len(payload):
1261
+ raise Corrupt("%s payload ends before all declared planes" % what)
1262
+ plane, consumed = huffman.decode_prefix(payload[offset:], element_count)
1263
+ offset += consumed
1264
+ planes.append(plane)
1265
+ if offset != len(payload):
1266
+ raise Corrupt(
1267
+ "%s payload has trailing bytes after the declared planes" % what
1268
+ )
1269
+ return planes
1270
+
1271
+
1272
+ def _delta_word_from_planes(dtype: str, planes, index: int) -> int:
1273
+ if dtype in ("bf16", "fp16"):
1274
+ return planes[0][index] | (planes[1][index] << 8)
1275
+ if dtype == "fp32":
1276
+ return (
1277
+ planes[0][index]
1278
+ | (planes[1][index] << 8)
1279
+ | (planes[2][index] << 16)
1280
+ | (planes[3][index] << 24)
1281
+ )
1282
+ if dtype == "q8_0":
1283
+ return planes[0][index]
1284
+ nibble = planes[0][index]
1285
+ if nibble > 15:
1286
+ raise Corrupt(
1287
+ "WEIGHT_DELTA q4_0 delta word %d is %d, outside 0..=15"
1288
+ % (index, nibble)
1289
+ )
1290
+ return nibble
1291
+
1292
+
1293
+ def _build_delta_bytes(dtype: str, planes, count: int) -> bytes:
1294
+ out = bytearray()
1295
+ for index in range(count):
1296
+ _push_word(dtype, _delta_word_from_planes(dtype, planes, index), out)
1297
+ return bytes(out)
1298
+
1299
+
1300
+ def _mod_mask(dtype: str) -> int:
1301
+ return (1 << WEIGHT_WORD_BITS[dtype]) - 1
1302
+
1303
+
1304
+ def _base_frame_descriptor(frame, limits: Limits) -> "tuple[str, list[int]]":
1305
+ from . import modes as _modes
1306
+
1307
+ mode = frame.header.mode
1308
+ if mode == 6:
1309
+ meta, _payload = _modes._parse_tensor_body(frame, limits)
1310
+ mapped = {
1311
+ "bf16": "bf16",
1312
+ "f16": "fp16",
1313
+ "f32": "fp32",
1314
+ "i8": "q8_0",
1315
+ "u8": "q4_0",
1316
+ }.get(meta["dtype"])
1317
+ if mapped is None:
1318
+ raise Malformed(
1319
+ "base frame dtype `%s` cannot back a WEIGHT_DELTA word plane"
1320
+ % meta["dtype"]
1321
+ )
1322
+ return mapped, meta["shape"]
1323
+ if mode == 9:
1324
+ meta = parse_float_descriptor(frame, limits)[0]
1325
+ mapped = {
1326
+ "bf16": "bf16",
1327
+ "fp16": "fp16",
1328
+ "fp32": "fp32",
1329
+ }.get(meta.dtype)
1330
+ if mapped is None:
1331
+ raise Malformed(
1332
+ "base frame dtype `%s` cannot back a WEIGHT_DELTA word plane"
1333
+ % meta.dtype
1334
+ )
1335
+ return mapped, meta.shape
1336
+ if mode == 11:
1337
+ meta = weight_metadata(frame, limits)
1338
+ return meta.dtype, meta.shape
1339
+ if mode == 14:
1340
+ meta = semantic_delta_metadata(frame, limits)
1341
+ return meta.dtype, meta.shape
1342
+ if mode == 16:
1343
+ meta = sparse_delta_metadata(frame, limits)
1344
+ return meta.dtype, meta.shape
1345
+ raise Malformed(
1346
+ "base frame mode `%s` does not carry tensor dtype/shape metadata"
1347
+ % frame.header.mode_name
1348
+ )
1349
+
1350
+
1351
+ def _base_bytes(
1352
+ frame,
1353
+ meta_base: bytes,
1354
+ resolver,
1355
+ limits: Limits,
1356
+ base_frame,
1357
+ what: str,
1358
+ ) -> "tuple[bytes, bytes | None]":
1359
+ from .modes import decode_bytes as _decode_bytes
1360
+
1361
+ if base_frame is not None:
1362
+ return _decode_bytes(base_frame, resolver, limits), base_frame
1363
+ if resolver is None:
1364
+ raise DependencyMissing(
1365
+ "frame needs dependency %s but no resolver was supplied"
1366
+ % meta_base.hex()
1367
+ )
1368
+ fetched = resolver.resolve_frame(meta_base)
1369
+ data = resolver.resolve(meta_base, frame.header.decoded_len, 0)
1370
+ if not isinstance(data, (bytes, bytearray, memoryview)):
1371
+ raise DependencyMismatch("resolver returned a non-bytes value")
1372
+ data = bytes(data)
1373
+ if len(data) != frame.header.decoded_len:
1374
+ raise DependencyMismatch(
1375
+ "dependency length %d differs from declared %d"
1376
+ % (len(data), frame.header.decoded_len)
1377
+ )
1378
+ from .hashing import content_root
1379
+
1380
+ if content_root(data) != meta_base:
1381
+ raise DependencyMismatch("dependency content hash mismatch")
1382
+ return data, fetched
1383
+
1384
+
1385
+ def _verify_base_frame(
1386
+ supplied: bytes, meta, limits: Limits, what: str
1387
+ ) -> None:
1388
+ from .frame import parse as _parse
1389
+
1390
+ parsed = _parse(supplied, limits)
1391
+ if parsed.header.content_hash != meta.base:
1392
+ raise DependencyMismatch(
1393
+ "%s supplied base frame content hash does not equal the descriptor base"
1394
+ % what
1395
+ )
1396
+ try:
1397
+ base_dtype, base_shape = _base_frame_descriptor(parsed, limits)
1398
+ except (Malformed, LimitExceeded) as exc:
1399
+ raise DependencyMismatch(exc.message)
1400
+ if base_dtype != meta.dtype:
1401
+ raise DependencyMismatch(
1402
+ "%s base frame dtype `%s` does not match %s dtype `%s`"
1403
+ % (what, base_dtype, what, meta.dtype)
1404
+ )
1405
+ if base_shape != meta.shape:
1406
+ raise DependencyMismatch(
1407
+ "%s base frame shape does not match the %s shape" % (what, what)
1408
+ )
1409
+
1410
+
1411
+ def _check_quantized(dtype: str, data: bytes, what: str) -> None:
1412
+ if dtype == "q4_0":
1413
+ for index, byte in enumerate(data):
1414
+ if byte > 15:
1415
+ raise Corrupt(
1416
+ "%s q4_0 word %d is %d, outside 0..=15"
1417
+ % (what, index, byte)
1418
+ )
1419
+
1420
+
1421
+ def decode_weight_delta(frame, resolver, limits: Limits, base_frame=None) -> bytes:
1422
+ descriptor_bytes, payload = _split_prefix(frame.body)
1423
+ meta = parse_weight_descriptor(descriptor_bytes, limits)
1424
+ count = validate_weight_meta(meta, limits)
1425
+ expected_len = _checked_mul(
1426
+ count, WEIGHT_STORAGE_BYTES[meta.dtype],
1427
+ "WEIGHT_DELTA element byte length",
1428
+ )
1429
+ if frame.header.decoded_len != expected_len:
1430
+ raise Corrupt(
1431
+ "WEIGHT_DELTA decoded_len is %d but the descriptor requires %d"
1432
+ % (frame.header.decoded_len, expected_len)
1433
+ )
1434
+ _check_len(expected_len, limits.max_decoded_bytes, "WEIGHT_DELTA element bytes")
1435
+ if meta.payload_len != len(payload):
1436
+ raise Malformed(
1437
+ "WEIGHT_DELTA payload is %d bytes but payload_len declares %d"
1438
+ % (len(payload), meta.payload_len)
1439
+ )
1440
+ if frame.header.dep_count != 1:
1441
+ raise Malformed("WEIGHT_DELTA requires exactly one dependency")
1442
+ dep = frame.dependencies[0]
1443
+ if dep.hash != meta.base:
1444
+ raise Malformed(
1445
+ "WEIGHT_DELTA dependency hash must equal the descriptor base"
1446
+ )
1447
+ if dep.decoded_len != frame.header.decoded_len:
1448
+ raise Malformed(
1449
+ "WEIGHT_DELTA dependency decoded_len must equal the frame decoded_len"
1450
+ )
1451
+ planes = _decode_planes(meta.dtype, payload, count, "WEIGHT_DELTA")
1452
+ delta_bytes = _build_delta_bytes(meta.dtype, planes, count)
1453
+ if values_hash(delta_bytes) != meta.delta_root:
1454
+ raise Corrupt("WEIGHT_DELTA delta_root mismatch")
1455
+
1456
+ base_bytes, supplied = _base_bytes(
1457
+ frame, meta.base, resolver, limits, base_frame, "WEIGHT_DELTA"
1458
+ )
1459
+ if supplied is not None:
1460
+ _verify_base_frame(supplied, meta, limits, "WEIGHT_DELTA")
1461
+ if len(base_bytes) != expected_len:
1462
+ raise DependencyMismatch(
1463
+ "base frame decodes to %d bytes but the descriptor requires %d"
1464
+ % (len(base_bytes), expected_len)
1465
+ )
1466
+ _check_quantized(meta.dtype, base_bytes, "WEIGHT_DELTA")
1467
+
1468
+ mask = _mod_mask(meta.dtype)
1469
+ target = bytearray()
1470
+ for index in range(count):
1471
+ delta = _delta_word_from_planes(meta.dtype, planes, index)
1472
+ base_word = _word_at(meta.dtype, base_bytes, index)
1473
+ _push_word(meta.dtype, (delta + base_word) & mask, target)
1474
+ if len(target) != expected_len:
1475
+ raise Corrupt(
1476
+ "WEIGHT_DELTA reconstructed element length is inconsistent"
1477
+ )
1478
+ if hash_domain("weft/chunk/v1", bytes(target)) != frame.header.content_hash:
1479
+ raise Corrupt(
1480
+ "WEIGHT_DELTA content_hash mismatch against the reconstructed tensor"
1481
+ )
1482
+ return bytes(target)
1483
+
1484
+
1485
+ # --- mode 14 ---------------------------------------------------------------
1486
+
1487
+
1488
+ @dataclass
1489
+ class SemanticDeltaMeta:
1490
+ dtype: str
1491
+ shape: "list[int]"
1492
+ element_count: int
1493
+ word_bits: int
1494
+ base: bytes
1495
+ base_dtype: str
1496
+ base_shape: "list[int]"
1497
+ base_semantic_root: bytes
1498
+ policy: SemanticPolicy
1499
+ delta_root: bytes
1500
+ plane_codec: str
1501
+ payload_len: int
1502
+
1503
+
1504
+ _SEMANTIC_DELTA_KEYS = (
1505
+ "schema",
1506
+ "dtype",
1507
+ "shape",
1508
+ "element_count",
1509
+ "word_bits",
1510
+ "base",
1511
+ "base_dtype",
1512
+ "base_shape",
1513
+ "base_semantic_root",
1514
+ "policy",
1515
+ "delta_root",
1516
+ "plane_codec",
1517
+ "payload_len",
1518
+ )
1519
+
1520
+
1521
+ def parse_semantic_delta_descriptor(descriptor_bytes: bytes, limits: Limits):
1522
+ map_value = _map(_decode_descriptor(descriptor_bytes, limits), "SEMANTIC_DELTA")
1523
+ _exact_keys(map_value, _SEMANTIC_DELTA_KEYS, "SEMANTIC_DELTA")
1524
+ schema = _text(map_value, "schema")
1525
+ if schema != SEMANTIC_DELTA_SCHEMA:
1526
+ raise Malformed("unsupported SEMANTIC_DELTA schema `%s`" % schema)
1527
+ dtype = _weight_dtype(_text(map_value, "dtype"))
1528
+ shape = _shape(map_value)
1529
+ element_count = _u64(map_value, "element_count")
1530
+ word_bits = _u64(map_value, "word_bits")
1531
+ base = _hash_field(map_value, "base")
1532
+ base_dtype = _weight_dtype(_text(map_value, "base_dtype"))
1533
+ base_shape = _array_u64(map_value, "base_shape")
1534
+ base_semantic_root = _hash_field(map_value, "base_semantic_root")
1535
+ policy = parse_semantic_policy(map_value.get("policy"), dtype, element_count)
1536
+ delta_root = _hash_field(map_value, "delta_root")
1537
+ plane_codec = _text(map_value, "plane_codec")
1538
+ payload_len = _u64(map_value, "payload_len")
1539
+ meta = SemanticDeltaMeta(
1540
+ dtype=dtype,
1541
+ shape=shape,
1542
+ element_count=element_count,
1543
+ word_bits=word_bits,
1544
+ base=base,
1545
+ base_dtype=base_dtype,
1546
+ base_shape=base_shape,
1547
+ base_semantic_root=base_semantic_root,
1548
+ policy=policy,
1549
+ delta_root=delta_root,
1550
+ plane_codec=plane_codec,
1551
+ payload_len=payload_len,
1552
+ )
1553
+ validate_semantic_delta_meta(meta, limits)
1554
+ return meta
1555
+
1556
+
1557
+ def validate_semantic_delta_meta(meta: SemanticDeltaMeta, limits: Limits) -> int:
1558
+ if meta.plane_codec != PLANE_CODEC:
1559
+ raise Malformed(
1560
+ "SEMANTIC_DELTA plane_codec `%s` is not implemented (only %s)"
1561
+ % (meta.plane_codec, PLANE_CODEC)
1562
+ )
1563
+ if WEIGHT_WORD_BITS[meta.dtype] != meta.word_bits:
1564
+ raise Malformed(
1565
+ "SEMANTIC_DELTA word_bits %d disagrees with dtype %s (%d bits)"
1566
+ % (meta.word_bits, meta.dtype, WEIGHT_WORD_BITS[meta.dtype])
1567
+ )
1568
+ if meta.base_dtype != meta.dtype:
1569
+ raise Malformed(
1570
+ "SEMANTIC_DELTA base_dtype `%s` must equal dtype `%s`"
1571
+ % (meta.base_dtype, meta.dtype)
1572
+ )
1573
+ if meta.base_shape != meta.shape:
1574
+ raise Malformed("SEMANTIC_DELTA base_shape must equal shape")
1575
+ count = _validate_shape(meta.shape, limits, "SEMANTIC_DELTA")
1576
+ if meta.element_count != count:
1577
+ raise Malformed(
1578
+ "SEMANTIC_DELTA element_count %d does not match shape product %d"
1579
+ % (meta.element_count, count)
1580
+ )
1581
+ _checked_mul(
1582
+ count, WEIGHT_STORAGE_BYTES[meta.dtype],
1583
+ "SEMANTIC_DELTA element byte length",
1584
+ )
1585
+ meta.policy.validate(meta.dtype, count)
1586
+ return count
1587
+
1588
+
1589
+ def semantic_delta_metadata(frame, limits: Limits) -> SemanticDeltaMeta:
1590
+ descriptor_bytes, _payload = _split_prefix(frame.body)
1591
+ return parse_semantic_delta_descriptor(descriptor_bytes, limits)
1592
+
1593
+
1594
+ def decode_semantic_delta(frame, resolver, limits: Limits, base_frame=None) -> bytes:
1595
+ descriptor_bytes, payload = _split_prefix(frame.body)
1596
+ meta = parse_semantic_delta_descriptor(descriptor_bytes, limits)
1597
+ count = validate_semantic_delta_meta(meta, limits)
1598
+ expected_len = _checked_mul(
1599
+ count, WEIGHT_STORAGE_BYTES[meta.dtype],
1600
+ "SEMANTIC_DELTA element byte length",
1601
+ )
1602
+ if frame.header.decoded_len != expected_len:
1603
+ raise Corrupt(
1604
+ "SEMANTIC_DELTA decoded_len is %d but the descriptor requires %d"
1605
+ % (frame.header.decoded_len, expected_len)
1606
+ )
1607
+ _check_len(expected_len, limits.max_decoded_bytes, "SEMANTIC_DELTA element bytes")
1608
+ if meta.payload_len != len(payload):
1609
+ raise Malformed(
1610
+ "SEMANTIC_DELTA payload is %d bytes but payload_len declares %d"
1611
+ % (len(payload), meta.payload_len)
1612
+ )
1613
+ if frame.header.dep_count != 1:
1614
+ raise Malformed("SEMANTIC_DELTA requires exactly one dependency")
1615
+ dep = frame.dependencies[0]
1616
+ if dep.hash != meta.base:
1617
+ raise Malformed(
1618
+ "SEMANTIC_DELTA dependency hash must equal the descriptor base"
1619
+ )
1620
+ if dep.decoded_len != frame.header.decoded_len:
1621
+ raise Malformed(
1622
+ "SEMANTIC_DELTA dependency decoded_len must equal the frame decoded_len"
1623
+ )
1624
+ planes = _decode_planes(meta.dtype, payload, count, "SEMANTIC_DELTA")
1625
+ delta_bytes = _build_delta_bytes(meta.dtype, planes, count)
1626
+ if values_hash(delta_bytes) != meta.delta_root:
1627
+ raise Corrupt("SEMANTIC_DELTA delta_root mismatch")
1628
+
1629
+ base_bytes, supplied = _base_bytes(
1630
+ frame, meta.base, resolver, limits, base_frame, "SEMANTIC_DELTA"
1631
+ )
1632
+ if supplied is not None:
1633
+ _verify_base_frame(supplied, meta, limits, "SEMANTIC_DELTA")
1634
+ if len(base_bytes) != expected_len:
1635
+ raise DependencyMismatch(
1636
+ "base frame decodes to %d bytes but the descriptor requires %d"
1637
+ % (len(base_bytes), expected_len)
1638
+ )
1639
+ _check_quantized(meta.dtype, base_bytes, "SEMANTIC_DELTA")
1640
+
1641
+ try:
1642
+ recomputed = semantic_root(
1643
+ meta.dtype, meta.shape, base_bytes, meta.policy, limits
1644
+ )
1645
+ except (Malformed, NonCanonical, LimitExceeded, Corrupt) as exc:
1646
+ raise DependencyMismatch(exc.message)
1647
+ if recomputed != meta.base_semantic_root:
1648
+ raise DependencyMismatch(
1649
+ "base semantic root does not equal the descriptor (no approximate "
1650
+ "reconstruction exists)"
1651
+ )
1652
+
1653
+ mask = _mod_mask(meta.dtype)
1654
+ target = bytearray()
1655
+ for index in range(count):
1656
+ delta = _delta_word_from_planes(meta.dtype, planes, index)
1657
+ base_word = _word_at(meta.dtype, base_bytes, index)
1658
+ _push_word(meta.dtype, (delta + base_word) & mask, target)
1659
+ if len(target) != expected_len:
1660
+ raise Corrupt(
1661
+ "SEMANTIC_DELTA reconstructed element length is inconsistent"
1662
+ )
1663
+ if hash_domain("weft/chunk/v1", bytes(target)) != frame.header.content_hash:
1664
+ raise Corrupt(
1665
+ "SEMANTIC_DELTA content_hash mismatch against the reconstructed tensor"
1666
+ )
1667
+ return bytes(target)
1668
+
1669
+
1670
+ # --- mode 16 ---------------------------------------------------------------
1671
+
1672
+ _SPARSE_KEYS = (
1673
+ "schema",
1674
+ "dtype",
1675
+ "shape",
1676
+ "element_count",
1677
+ "storage_bytes",
1678
+ "op",
1679
+ "sparse_codec",
1680
+ "base_content_root",
1681
+ "delta_root",
1682
+ )
1683
+
1684
+
1685
+ @dataclass
1686
+ class SparseDeltaMeta:
1687
+ dtype: str
1688
+ shape: "list[int]"
1689
+ element_count: int
1690
+ storage_bytes: int
1691
+ op: str
1692
+ sparse_codec: str
1693
+ base: bytes
1694
+ delta_root: bytes
1695
+
1696
+
1697
+ def parse_sparse_descriptor(descriptor_bytes: bytes, limits: Limits) -> SparseDeltaMeta:
1698
+ map_value = _map(_decode_descriptor(descriptor_bytes, limits), "SPARSE_DELTA")
1699
+ _exact_keys(map_value, _SPARSE_KEYS, "SPARSE_DELTA")
1700
+ schema = _text(map_value, "schema")
1701
+ if schema != SPARSE_DELTA_SCHEMA:
1702
+ raise Malformed("unsupported SPARSE_DELTA schema `%s`" % schema)
1703
+ dtype = _weight_dtype(_text(map_value, "dtype"))
1704
+ shape = _shape(map_value)
1705
+ element_count = _u64(map_value, "element_count")
1706
+ storage_bytes = _u64(map_value, "storage_bytes")
1707
+ op = _text(map_value, "op")
1708
+ sparse_codec = _text(map_value, "sparse_codec")
1709
+ base = _hash_field(map_value, "base_content_root")
1710
+ delta_root = _hash_field(map_value, "delta_root")
1711
+ meta = SparseDeltaMeta(
1712
+ dtype=dtype,
1713
+ shape=shape,
1714
+ element_count=element_count,
1715
+ storage_bytes=storage_bytes,
1716
+ op=op,
1717
+ sparse_codec=sparse_codec,
1718
+ base=base,
1719
+ delta_root=delta_root,
1720
+ )
1721
+ validate_sparse_meta(meta, limits)
1722
+ return meta
1723
+
1724
+
1725
+ def validate_sparse_meta(meta: SparseDeltaMeta, limits: Limits) -> int:
1726
+ if meta.op != DELTA_OP:
1727
+ raise Malformed(
1728
+ "SPARSE_DELTA op `%s` is not implemented (only %s)"
1729
+ % (meta.op, DELTA_OP)
1730
+ )
1731
+ if meta.sparse_codec != SPARSE_CODEC:
1732
+ raise Malformed(
1733
+ "SPARSE_DELTA sparse_codec `%s` is not implemented (only %s)"
1734
+ % (meta.sparse_codec, SPARSE_CODEC)
1735
+ )
1736
+ if meta.storage_bytes != WEIGHT_STORAGE_BYTES[meta.dtype]:
1737
+ raise Malformed(
1738
+ "SPARSE_DELTA storage_bytes %d disagrees with dtype %s (%d bytes)"
1739
+ % (
1740
+ meta.storage_bytes,
1741
+ meta.dtype,
1742
+ WEIGHT_STORAGE_BYTES[meta.dtype],
1743
+ )
1744
+ )
1745
+ count = _validate_shape(meta.shape, limits, "SPARSE_DELTA")
1746
+ if meta.element_count != count:
1747
+ raise Malformed(
1748
+ "SPARSE_DELTA element_count %d does not match shape product %d"
1749
+ % (meta.element_count, count)
1750
+ )
1751
+ _checked_mul(
1752
+ count, WEIGHT_STORAGE_BYTES[meta.dtype],
1753
+ "SPARSE_DELTA element byte length",
1754
+ )
1755
+ return count
1756
+
1757
+
1758
+ def sparse_delta_metadata(frame, limits: Limits) -> SparseDeltaMeta:
1759
+ descriptor_bytes, _payload = _split_prefix(frame.body)
1760
+ return parse_sparse_descriptor(descriptor_bytes, limits)
1761
+
1762
+
1763
+ def _decode_runs(
1764
+ dtype: str, payload: bytes, element_count: int, limits: Limits
1765
+ ) -> bytes:
1766
+ if len(payload) < 4:
1767
+ raise Malformed("SPARSE_DELTA body is shorter than its 4-byte run_count")
1768
+ run_count = int.from_bytes(payload[0:4], "little")
1769
+ if run_count > limits.max_sparse_runs:
1770
+ raise LimitExceeded(
1771
+ "SPARSE_DELTA run_count %d exceeds limit %d"
1772
+ % (run_count, limits.max_sparse_runs)
1773
+ )
1774
+ if run_count > element_count:
1775
+ raise Malformed(
1776
+ "SPARSE_DELTA run_count %d exceeds element_count %d"
1777
+ % (run_count, element_count)
1778
+ )
1779
+ storage = WEIGHT_STORAGE_BYTES[dtype]
1780
+ delta_len = _checked_mul(element_count, storage, "SPARSE_DELTA delta bytes")
1781
+ delta = bytearray(delta_len)
1782
+ position = 4
1783
+ emitted = 0
1784
+ for record in range(run_count):
1785
+ zero_run, position = varint.decode(payload, position)
1786
+ emitted = _checked_add(
1787
+ emitted, zero_run, "SPARSE_DELTA zero-run length"
1788
+ )
1789
+ if emitted >= element_count:
1790
+ raise Malformed(
1791
+ "SPARSE_DELTA run %d skips %d words past element_count %d"
1792
+ % (record, zero_run, element_count)
1793
+ )
1794
+ end = position + storage
1795
+ if end > len(payload):
1796
+ raise Malformed("SPARSE_DELTA run %d record is truncated" % record)
1797
+ word = payload[position:end]
1798
+ if not any(word):
1799
+ raise Corrupt(
1800
+ "SPARSE_DELTA run %d stores a zero word (non-canonical: zeros "
1801
+ "are implicit)" % record
1802
+ )
1803
+ if dtype == "q4_0" and word[0] > 15:
1804
+ raise Corrupt(
1805
+ "SPARSE_DELTA q4_0 run %d word is %d, outside 0..=15"
1806
+ % (record, word[0])
1807
+ )
1808
+ offset = emitted * storage
1809
+ delta[offset : offset + storage] = word
1810
+ position = end
1811
+ emitted += 1
1812
+ if position != len(payload):
1813
+ raise Corrupt(
1814
+ "SPARSE_DELTA payload has %d trailing bytes after the %d declared runs"
1815
+ % (len(payload) - position, run_count)
1816
+ )
1817
+ return bytes(delta)
1818
+
1819
+
1820
+ def decode_sparse_delta(frame, resolver, limits: Limits, base_frame=None) -> bytes:
1821
+ descriptor_bytes, payload = _split_prefix(frame.body)
1822
+ meta = parse_sparse_descriptor(descriptor_bytes, limits)
1823
+ count = validate_sparse_meta(meta, limits)
1824
+ expected_len = _checked_mul(
1825
+ count, WEIGHT_STORAGE_BYTES[meta.dtype],
1826
+ "SPARSE_DELTA element byte length",
1827
+ )
1828
+ if frame.header.decoded_len != expected_len:
1829
+ raise Corrupt(
1830
+ "SPARSE_DELTA decoded_len is %d but the descriptor requires %d"
1831
+ % (frame.header.decoded_len, expected_len)
1832
+ )
1833
+ _check_len(expected_len, limits.max_decoded_bytes, "SPARSE_DELTA element bytes")
1834
+ if frame.header.dep_count != 1:
1835
+ raise Malformed("SPARSE_DELTA requires exactly one dependency")
1836
+ dep = frame.dependencies[0]
1837
+ if dep.hash != meta.base:
1838
+ raise Malformed(
1839
+ "SPARSE_DELTA dependency hash must equal the descriptor "
1840
+ "base_content_root"
1841
+ )
1842
+ if dep.decoded_len != frame.header.decoded_len:
1843
+ raise Malformed(
1844
+ "SPARSE_DELTA dependency decoded_len must equal the frame decoded_len"
1845
+ )
1846
+ delta_bytes = _decode_runs(meta.dtype, payload, count, limits)
1847
+ if values_hash(delta_bytes) != meta.delta_root:
1848
+ raise Corrupt("SPARSE_DELTA delta_root mismatch")
1849
+
1850
+ base_bytes, supplied = _base_bytes(
1851
+ frame, meta.base, resolver, limits, base_frame, "SPARSE_DELTA"
1852
+ )
1853
+ if supplied is not None:
1854
+ _verify_base_frame(supplied, meta, limits, "SPARSE_DELTA")
1855
+ if len(base_bytes) != expected_len:
1856
+ raise DependencyMismatch(
1857
+ "base frame decodes to %d bytes but the descriptor requires %d"
1858
+ % (len(base_bytes), expected_len)
1859
+ )
1860
+ _check_quantized(meta.dtype, base_bytes, "SPARSE_DELTA")
1861
+
1862
+ mask = _mod_mask(meta.dtype)
1863
+ target = bytearray()
1864
+ for index in range(count):
1865
+ delta = _word_at(meta.dtype, delta_bytes, index)
1866
+ base_word = _word_at(meta.dtype, base_bytes, index)
1867
+ _push_word(meta.dtype, (delta + base_word) & mask, target)
1868
+ if len(target) != expected_len:
1869
+ raise Corrupt(
1870
+ "SPARSE_DELTA reconstructed element length is inconsistent"
1871
+ )
1872
+ if hash_domain("weft/chunk/v1", bytes(target)) != frame.header.content_hash:
1873
+ raise Corrupt(
1874
+ "SPARSE_DELTA content_hash mismatch against the reconstructed tensor"
1875
+ )
1876
+ return bytes(target)
1877
+
1878
+
1879
+ def activation_stream_root(
1880
+ stream: str,
1881
+ step: int,
1882
+ rank: int,
1883
+ world: int,
1884
+ chunk_count: int,
1885
+ digest_roots: "list[bytes]",
1886
+ limits: Limits,
1887
+ ) -> bytes:
1888
+ if len(digest_roots) != chunk_count:
1889
+ raise Malformed(
1890
+ "ACTIVATION_STREAM stream root needs %d digests, got %d"
1891
+ % (chunk_count, len(digest_roots))
1892
+ )
1893
+ record = {
1894
+ "stream": stream,
1895
+ "step": wcb.U64(step),
1896
+ "rank": wcb.U64(rank),
1897
+ "world": wcb.U64(world),
1898
+ "chunk_count": wcb.U64(chunk_count),
1899
+ "digest_roots": [bytes(digest) for digest in digest_roots],
1900
+ }
1901
+ return hash_domain(DOMAIN_AI_STREAM, wcb.encode(record, limits))
1902
+
1903
+
1904
+ # ---------------------------------------------------------------------------
1905
+ # Mode 12 -- KV_CACHE
1906
+ # ---------------------------------------------------------------------------
1907
+
1908
+ _KV_KEYS = (
1909
+ "schema",
1910
+ "model",
1911
+ "layer",
1912
+ "session",
1913
+ "page_index",
1914
+ "page_count",
1915
+ "seq_start",
1916
+ "seq_end",
1917
+ "dtype",
1918
+ "head_count",
1919
+ "head_dim",
1920
+ "shape",
1921
+ "page_bytes",
1922
+ "plane_codec",
1923
+ "payload_len",
1924
+ )
1925
+
1926
+
1927
+ @dataclass
1928
+ class KvMeta:
1929
+ model: str
1930
+ layer: int
1931
+ session: str
1932
+ page_index: int
1933
+ page_count: int
1934
+ seq_start: int
1935
+ seq_end: int
1936
+ dtype: str
1937
+ head_count: int
1938
+ head_dim: int
1939
+ shape: "list[int]"
1940
+ page_bytes: int
1941
+ plane_codec: str
1942
+ payload_len: int
1943
+
1944
+
1945
+ def _kv_dtype(name: str) -> str:
1946
+ if name not in KV_WORD_BYTES:
1947
+ raise Malformed(
1948
+ "unknown KV_CACHE dtype `%s` (v1 allows bf16, fp16, fp32)" % name
1949
+ )
1950
+ return name
1951
+
1952
+
1953
+ def parse_kv_descriptor(descriptor_bytes: bytes, limits: Limits) -> KvMeta:
1954
+ map_value = _map(_decode_descriptor(descriptor_bytes, limits), "KV_CACHE")
1955
+ _exact_keys(map_value, _KV_KEYS, "KV_CACHE")
1956
+ schema = _text(map_value, "schema")
1957
+ if schema != KV_CACHE_SCHEMA:
1958
+ raise Malformed("unsupported KV_CACHE schema `%s`" % schema)
1959
+ meta = KvMeta(
1960
+ model=_text(map_value, "model"),
1961
+ layer=_u32(map_value, "layer"),
1962
+ session=_text(map_value, "session"),
1963
+ page_index=_u64(map_value, "page_index"),
1964
+ page_count=_u64(map_value, "page_count"),
1965
+ seq_start=_u64(map_value, "seq_start"),
1966
+ seq_end=_u64(map_value, "seq_end"),
1967
+ dtype=_kv_dtype(_text(map_value, "dtype")),
1968
+ head_count=_u32(map_value, "head_count"),
1969
+ head_dim=_u32(map_value, "head_dim"),
1970
+ shape=_shape(map_value),
1971
+ page_bytes=_u64(map_value, "page_bytes"),
1972
+ plane_codec=_text(map_value, "plane_codec"),
1973
+ payload_len=_u64(map_value, "payload_len"),
1974
+ )
1975
+ validate_kv_meta(meta, limits)
1976
+ return meta
1977
+
1978
+
1979
+ def validate_kv_meta(meta: KvMeta, limits: Limits) -> int:
1980
+ if meta.plane_codec != PLANE_CODEC:
1981
+ raise Malformed(
1982
+ "KV_CACHE plane_codec `%s` is not implemented (only %s)"
1983
+ % (meta.plane_codec, PLANE_CODEC)
1984
+ )
1985
+ if meta.page_count == 0:
1986
+ raise Malformed("KV_CACHE page_count must be at least 1")
1987
+ if meta.page_index >= meta.page_count:
1988
+ raise Malformed(
1989
+ "KV_CACHE page_index %d is not below page_count %d"
1990
+ % (meta.page_index, meta.page_count)
1991
+ )
1992
+ if meta.seq_start >= meta.seq_end:
1993
+ raise Malformed(
1994
+ "KV_CACHE seq_start %d must be less than seq_end %d"
1995
+ % (meta.seq_start, meta.seq_end)
1996
+ )
1997
+ rows = meta.seq_end - meta.seq_start
1998
+ if meta.head_count == 0:
1999
+ raise Malformed("KV_CACHE head_count must be greater than zero")
2000
+ if meta.head_dim == 0:
2001
+ raise Malformed("KV_CACHE head_dim must be greater than zero")
2002
+ expected_shape = [rows, meta.head_count, meta.head_dim]
2003
+ if meta.shape != expected_shape:
2004
+ raise Malformed(
2005
+ "KV_CACHE shape %r must equal [page_rows, head_count, head_dim] = %r"
2006
+ % (meta.shape, expected_shape)
2007
+ )
2008
+ if len(meta.shape) > limits.max_tensor_rank:
2009
+ raise LimitExceeded(
2010
+ "KV_CACHE rank %d exceeds limit %d"
2011
+ % (len(meta.shape), limits.max_tensor_rank)
2012
+ )
2013
+ element_count = _checked_mul(rows, meta.head_count, "KV_CACHE element count")
2014
+ element_count = _checked_mul(
2015
+ element_count, meta.head_dim, "KV_CACHE element count"
2016
+ )
2017
+ if element_count > limits.max_tensor_elements:
2018
+ raise LimitExceeded(
2019
+ "KV_CACHE element count %d exceeds limit %d"
2020
+ % (element_count, limits.max_tensor_elements)
2021
+ )
2022
+ computed = _checked_mul(
2023
+ element_count, KV_WORD_BYTES[meta.dtype], "KV_CACHE page byte count"
2024
+ )
2025
+ if meta.page_bytes != computed:
2026
+ raise Malformed(
2027
+ "KV_CACHE page_bytes %d does not equal page_rows * head_count * "
2028
+ "head_dim * word_bytes = %d" % (meta.page_bytes, computed)
2029
+ )
2030
+ _check_len(meta.page_bytes, limits.max_decoded_bytes, "KV_CACHE page")
2031
+ return computed
2032
+
2033
+
2034
+ def _decode_byte_planes(
2035
+ payload: bytes, word_bytes: int, byte_len: int, what: str
2036
+ ) -> bytes:
2037
+ if byte_len % word_bytes != 0:
2038
+ raise Malformed("%s byte length is not a whole number of elements" % what)
2039
+ element_count = byte_len // word_bytes
2040
+ plane_count = word_bytes
2041
+ planes = []
2042
+ offset = 0
2043
+ for _ in range(plane_count):
2044
+ if offset >= len(payload):
2045
+ raise Corrupt("%s payload ends before all declared planes" % what)
2046
+ plane, consumed = huffman.decode_prefix(payload[offset:], element_count)
2047
+ offset += consumed
2048
+ planes.append(plane)
2049
+ if offset != len(payload):
2050
+ raise Corrupt(
2051
+ "%s payload has trailing bytes after the declared planes" % what
2052
+ )
2053
+ out = bytearray()
2054
+ for index in range(element_count):
2055
+ word = 0
2056
+ for plane_index, plane in enumerate(planes):
2057
+ word |= plane[index] << (8 * plane_index)
2058
+ out.extend(word.to_bytes(word_bytes, "little"))
2059
+ return bytes(out)
2060
+
2061
+
2062
+ def decode_kv_cache(frame, limits: Limits) -> bytes:
2063
+ descriptor_bytes, payload = _split_prefix(frame.body)
2064
+ meta = parse_kv_descriptor(descriptor_bytes, limits)
2065
+ page_bytes = validate_kv_meta(meta, limits)
2066
+ if frame.header.decoded_len != page_bytes:
2067
+ raise Corrupt(
2068
+ "KV_CACHE decoded_len is %d but the descriptor requires %d"
2069
+ % (frame.header.decoded_len, page_bytes)
2070
+ )
2071
+ if meta.payload_len != len(payload):
2072
+ raise Malformed(
2073
+ "KV_CACHE payload is %d bytes but payload_len declares %d"
2074
+ % (len(payload), meta.payload_len)
2075
+ )
2076
+ out = _decode_byte_planes(
2077
+ payload, KV_WORD_BYTES[meta.dtype], page_bytes, "KV_CACHE"
2078
+ )
2079
+ if len(out) != page_bytes:
2080
+ raise Corrupt("KV_CACHE reconstructed page length is inconsistent")
2081
+ if hash_domain("weft/chunk/v1", out) != frame.header.content_hash:
2082
+ raise Corrupt(
2083
+ "KV_CACHE content_hash mismatch against the reconstructed page"
2084
+ )
2085
+ return out
2086
+
2087
+
2088
+ # ---------------------------------------------------------------------------
2089
+ # Mode 13 -- ACTIVATION_STREAM
2090
+ # ---------------------------------------------------------------------------
2091
+
2092
+ _ACTIVATION_KEYS = (
2093
+ "schema",
2094
+ "stream",
2095
+ "step",
2096
+ "rank",
2097
+ "world",
2098
+ "chunk_index",
2099
+ "chunk_count",
2100
+ "dtype",
2101
+ "shape",
2102
+ "element_count",
2103
+ "digest_root",
2104
+ "plane_codec",
2105
+ "payload_len",
2106
+ )
2107
+
2108
+
2109
+ @dataclass
2110
+ class ActivationMeta:
2111
+ stream: str
2112
+ step: int
2113
+ rank: int
2114
+ world: int
2115
+ chunk_index: int
2116
+ chunk_count: int
2117
+ dtype: str
2118
+ shape: "list[int]"
2119
+ element_count: int
2120
+ digest_root: bytes
2121
+ plane_codec: str
2122
+ payload_len: int
2123
+
2124
+
2125
+ def _stream_dtype(name: str) -> str:
2126
+ if name not in KV_WORD_BYTES:
2127
+ raise Malformed(
2128
+ "unknown ACTIVATION_STREAM dtype `%s` (v1 allows bf16, fp16, fp32)"
2129
+ % name
2130
+ )
2131
+ return name
2132
+
2133
+
2134
+ def _shape_product(shape: "list[int]", limits: Limits) -> int:
2135
+ if not shape:
2136
+ raise Malformed("ACTIVATION_STREAM shape rank must be at least 1")
2137
+ return _validate_shape(shape, limits, "ACTIVATION_STREAM")
2138
+
2139
+
2140
+ def parse_activation_descriptor(
2141
+ descriptor_bytes: bytes, limits: Limits
2142
+ ) -> ActivationMeta:
2143
+ map_value = _map(
2144
+ _decode_descriptor(descriptor_bytes, limits), "ACTIVATION_STREAM"
2145
+ )
2146
+ _exact_keys(map_value, _ACTIVATION_KEYS, "ACTIVATION_STREAM")
2147
+ schema = _text(map_value, "schema")
2148
+ if schema != ACTIVATION_STREAM_SCHEMA:
2149
+ raise Malformed("unsupported ACTIVATION_STREAM schema `%s`" % schema)
2150
+ stream = _text(map_value, "stream")
2151
+ if stream not in ("activation", "gradient"):
2152
+ raise Malformed(
2153
+ "unknown ACTIVATION_STREAM stream `%s` (v1 allows activation, "
2154
+ "gradient)" % stream
2155
+ )
2156
+ meta = ActivationMeta(
2157
+ stream=stream,
2158
+ step=_u64(map_value, "step"),
2159
+ rank=_u32(map_value, "rank"),
2160
+ world=_u32(map_value, "world"),
2161
+ chunk_index=_u64(map_value, "chunk_index"),
2162
+ chunk_count=_u64(map_value, "chunk_count"),
2163
+ dtype=_stream_dtype(_text(map_value, "dtype")),
2164
+ shape=_shape(map_value),
2165
+ element_count=_u64(map_value, "element_count"),
2166
+ digest_root=_hash_field(map_value, "digest_root"),
2167
+ plane_codec=_text(map_value, "plane_codec"),
2168
+ payload_len=_u64(map_value, "payload_len"),
2169
+ )
2170
+ validate_activation_meta(meta, limits)
2171
+ return meta
2172
+
2173
+
2174
+ def validate_activation_meta(meta: ActivationMeta, limits: Limits) -> int:
2175
+ if meta.plane_codec != PLANE_CODEC:
2176
+ raise Malformed(
2177
+ "ACTIVATION_STREAM plane_codec `%s` is not implemented (only %s)"
2178
+ % (meta.plane_codec, PLANE_CODEC)
2179
+ )
2180
+ if meta.world == 0:
2181
+ raise Malformed("ACTIVATION_STREAM world must be at least 1")
2182
+ if meta.rank >= meta.world:
2183
+ raise Malformed(
2184
+ "ACTIVATION_STREAM rank %d must be below world %d"
2185
+ % (meta.rank, meta.world)
2186
+ )
2187
+ if meta.chunk_count == 0:
2188
+ raise Malformed("ACTIVATION_STREAM chunk_count must be at least 1")
2189
+ product = _shape_product(meta.shape, limits)
2190
+ byte_len = _checked_mul(
2191
+ product, KV_WORD_BYTES[meta.dtype], "ACTIVATION_STREAM chunk byte count"
2192
+ )
2193
+ _check_len(byte_len, limits.max_decoded_bytes, "ACTIVATION_STREAM chunk")
2194
+ if meta.chunk_index >= meta.chunk_count:
2195
+ raise Malformed(
2196
+ "ACTIVATION_STREAM chunk_index %d is not below chunk_count %d"
2197
+ % (meta.chunk_index, meta.chunk_count)
2198
+ )
2199
+ if meta.element_count != product:
2200
+ raise Malformed(
2201
+ "ACTIVATION_STREAM element_count %d does not match shape product %d"
2202
+ % (meta.element_count, product)
2203
+ )
2204
+ return byte_len
2205
+
2206
+
2207
+ def activation_metadata(frame, limits: Limits) -> ActivationMeta:
2208
+ descriptor_bytes, _payload = _split_prefix(frame.body)
2209
+ return parse_activation_descriptor(descriptor_bytes, limits)
2210
+
2211
+
2212
+ def decode_activation_stream(frame, limits: Limits) -> bytes:
2213
+ descriptor_bytes, payload = _split_prefix(frame.body)
2214
+ meta = parse_activation_descriptor(descriptor_bytes, limits)
2215
+ byte_len = validate_activation_meta(meta, limits)
2216
+ if frame.header.decoded_len != byte_len:
2217
+ raise Corrupt(
2218
+ "ACTIVATION_STREAM decoded_len is %d but the descriptor requires %d"
2219
+ % (frame.header.decoded_len, byte_len)
2220
+ )
2221
+ if meta.payload_len != len(payload):
2222
+ raise Malformed(
2223
+ "ACTIVATION_STREAM payload is %d bytes but payload_len declares %d"
2224
+ % (len(payload), meta.payload_len)
2225
+ )
2226
+ out = _decode_byte_planes(
2227
+ payload, KV_WORD_BYTES[meta.dtype], byte_len, "ACTIVATION_STREAM"
2228
+ )
2229
+ if len(out) != byte_len:
2230
+ raise Corrupt(
2231
+ "ACTIVATION_STREAM reconstructed chunk length is inconsistent"
2232
+ )
2233
+ if values_hash(out) != meta.digest_root:
2234
+ raise Corrupt("ACTIVATION_STREAM digest_root mismatch")
2235
+ if hash_domain("weft/chunk/v1", out) != frame.header.content_hash:
2236
+ raise Corrupt(
2237
+ "ACTIVATION_STREAM content_hash mismatch against the reconstructed chunk"
2238
+ )
2239
+ return out
2240
+
2241
+
2242
+ # ---------------------------------------------------------------------------
2243
+ # Mode 15 -- TENSOR_TILED
2244
+ # ---------------------------------------------------------------------------
2245
+
2246
+ _TILED_KEYS = (
2247
+ "schema",
2248
+ "dtype",
2249
+ "shape",
2250
+ "element_count",
2251
+ "layout",
2252
+ "tile_shape",
2253
+ "tile_counts",
2254
+ "tile_codec",
2255
+ "tiles",
2256
+ "payload_len",
2257
+ "logical_root",
2258
+ )
2259
+
2260
+ _TILE_KEYS = (
2261
+ "index",
2262
+ "dims",
2263
+ "offset",
2264
+ "compressed_len",
2265
+ "element_len",
2266
+ "tile_root",
2267
+ )
2268
+
2269
+
2270
+ @dataclass
2271
+ class TileEntry:
2272
+ index: int
2273
+ dims: "list[int]"
2274
+ offset: int
2275
+ compressed_len: int
2276
+ element_len: int
2277
+ tile_root: bytes
2278
+
2279
+
2280
+ @dataclass
2281
+ class TiledMeta:
2282
+ dtype: str
2283
+ shape: "list[int]"
2284
+ element_count: int
2285
+ tile_shape: "list[int]"
2286
+ tile_counts: "list[int]"
2287
+ tiles: "list[TileEntry]"
2288
+ payload_len: int
2289
+ logical_root: bytes
2290
+
2291
+
2292
+ def _tensor_dtype(name: str) -> str:
2293
+ if name not in TENSOR_DTYPE_SIZES:
2294
+ raise Malformed(
2295
+ "unknown tensor dtype `%s` (v1 allows f16, bf16, f32, f64, i8, u8, "
2296
+ "i16, i32, i64)" % name
2297
+ )
2298
+ return name
2299
+
2300
+
2301
+ def _tile_counts(shape: "list[int]", tile_shape: "list[int]") -> "list[int]":
2302
+ if len(shape) != len(tile_shape):
2303
+ raise Malformed(
2304
+ "tensor tile rank %d does not equal shape rank %d"
2305
+ % (len(tile_shape), len(shape))
2306
+ )
2307
+ counts = []
2308
+ for axis, (extent, tile) in enumerate(zip(shape, tile_shape)):
2309
+ if tile == 0:
2310
+ raise Malformed(
2311
+ "tensor tile shape axis %d must be greater than zero" % axis
2312
+ )
2313
+ if extent % tile != 0:
2314
+ raise Malformed(
2315
+ "tensor tile shape axis %d extent %d does not divide shape "
2316
+ "extent %d" % (axis, tile, extent)
2317
+ )
2318
+ counts.append(extent // tile)
2319
+ return counts
2320
+
2321
+
2322
+ def parse_tiled_descriptor(
2323
+ descriptor_bytes: bytes, decoded_len: int, payload_len: int, limits: Limits
2324
+ ) -> TiledMeta:
2325
+ map_value = _map(_decode_descriptor(descriptor_bytes, limits), "tiled tensor")
2326
+ _exact_keys(map_value, _TILED_KEYS, "tiled tensor")
2327
+ schema = _text(map_value, "schema")
2328
+ if schema != TENSOR_TILED_SCHEMA:
2329
+ raise Malformed("unsupported tiled tensor schema `%s`" % schema)
2330
+ dtype = _tensor_dtype(_text(map_value, "dtype"))
2331
+ layout = _text(map_value, "layout")
2332
+ if layout != TILE_LAYOUT:
2333
+ raise Malformed(
2334
+ "tiled tensor layout `%s` is specified but not implemented in v1 "
2335
+ "(only %s)" % (layout, TILE_LAYOUT)
2336
+ )
2337
+ tile_codec = _text(map_value, "tile_codec")
2338
+ if tile_codec != PLANE_CODEC:
2339
+ raise Malformed(
2340
+ "unsupported tensor tile codec `%s` (v1 allows `%s`)"
2341
+ % (tile_codec, PLANE_CODEC)
2342
+ )
2343
+ shape = _shape(map_value)
2344
+ element_count = _validate_shape(shape, limits, "tiled tensor")
2345
+ declared_count = _u64(map_value, "element_count")
2346
+ if declared_count != element_count:
2347
+ raise Malformed(
2348
+ "tiled tensor element_count %d does not match shape product %d"
2349
+ % (declared_count, element_count)
2350
+ )
2351
+ element_size = TENSOR_DTYPE_SIZES[dtype]
2352
+ decoded_bytes = _checked_mul(
2353
+ element_count, element_size, "tiled tensor decoded bytes"
2354
+ )
2355
+ _check_len(decoded_bytes, limits.max_decoded_bytes, "tiled tensor decoded bytes")
2356
+ if decoded_len != decoded_bytes:
2357
+ raise Corrupt(
2358
+ "tiled tensor decoded_len is %d but the descriptor requires %d"
2359
+ % (decoded_len, decoded_bytes)
2360
+ )
2361
+ tile_shape = _array_u64(map_value, "tile_shape")
2362
+ counts = _tile_counts(shape, tile_shape)
2363
+ declared_counts = _array_u64(map_value, "tile_counts")
2364
+ if declared_counts != counts:
2365
+ raise Malformed(
2366
+ "tiled tensor tile_counts %r do not match shape/tile_shape %r"
2367
+ % (declared_counts, counts)
2368
+ )
2369
+ tile_elements = _validate_shape(tile_shape, limits, "tensor tile")
2370
+ tile_element_len = _checked_mul(
2371
+ tile_elements, element_size, "tensor tile byte length"
2372
+ )
2373
+ tile_count = 1
2374
+ for count in counts:
2375
+ tile_count = _checked_mul(tile_count, count, "tensor tile count")
2376
+ if tile_count > limits.max_structured_items:
2377
+ raise LimitExceeded(
2378
+ "tensor tile count %d exceeds limit %d"
2379
+ % (tile_count, limits.max_structured_items)
2380
+ )
2381
+ declared_payload_len = _u64(map_value, "payload_len")
2382
+ if declared_payload_len != payload_len:
2383
+ raise Malformed(
2384
+ "tiled tensor payload is %d bytes but metadata declares %d"
2385
+ % (payload_len, declared_payload_len)
2386
+ )
2387
+ entries = map_value.get("tiles")
2388
+ if not isinstance(entries, list):
2389
+ raise Malformed("tiled tensor `tiles` must be an array")
2390
+ if len(entries) != tile_count:
2391
+ raise Malformed(
2392
+ "tiled tensor has %d tile entries but the grid requires %d"
2393
+ % (len(entries), tile_count)
2394
+ )
2395
+ tiles = []
2396
+ expected_offset = 0
2397
+ for position, item in enumerate(entries):
2398
+ entry = _map(item, "tensor tile")
2399
+ _exact_keys(entry, _TILE_KEYS, "tensor tile")
2400
+ index = _u64(entry, "index")
2401
+ if index != position:
2402
+ raise Malformed(
2403
+ "tiled tensor entry %d declares index %d; entries must be in "
2404
+ "canonical order" % (position, index)
2405
+ )
2406
+ dims = _array_u64(entry, "dims")
2407
+ if dims != tile_shape:
2408
+ raise Malformed(
2409
+ "tensor tile %d dims %r do not equal tile_shape %r; partial "
2410
+ "tiles are not defined in v1" % (position, dims, tile_shape)
2411
+ )
2412
+ declared_element_len = _u64(entry, "element_len")
2413
+ if declared_element_len != tile_element_len:
2414
+ raise Malformed(
2415
+ "tensor tile %d element_len %d does not match %d"
2416
+ % (position, declared_element_len, tile_element_len)
2417
+ )
2418
+ offset = _u64(entry, "offset")
2419
+ if offset != expected_offset:
2420
+ raise Malformed(
2421
+ "tensor tile %d offset %d is not the contiguous offset %d"
2422
+ % (position, offset, expected_offset)
2423
+ )
2424
+ compressed_len = _u64(entry, "compressed_len")
2425
+ if compressed_len == 0:
2426
+ raise Malformed(
2427
+ "tensor tile %d has an empty stored range" % position
2428
+ )
2429
+ expected_offset = _checked_add(
2430
+ expected_offset, compressed_len, "tensor tile offset"
2431
+ )
2432
+ root = _bytes_of(entry, "tile_root")
2433
+ if len(root) != HASH_LEN:
2434
+ raise Malformed("tensor tile_root must be 32 bytes")
2435
+ tiles.append(
2436
+ TileEntry(
2437
+ index=index,
2438
+ dims=dims,
2439
+ offset=offset,
2440
+ compressed_len=compressed_len,
2441
+ element_len=declared_element_len,
2442
+ tile_root=root,
2443
+ )
2444
+ )
2445
+ if expected_offset != declared_payload_len:
2446
+ raise Malformed(
2447
+ "tensor tile ranges cover %d bytes but payload_len is %d"
2448
+ % (expected_offset, declared_payload_len)
2449
+ )
2450
+ logical_root = _bytes_of(map_value, "logical_root")
2451
+ if len(logical_root) != HASH_LEN:
2452
+ raise Malformed("tiled tensor logical_root must be 32 bytes")
2453
+ return TiledMeta(
2454
+ dtype=dtype,
2455
+ shape=shape,
2456
+ element_count=element_count,
2457
+ tile_shape=tile_shape,
2458
+ tile_counts=counts,
2459
+ tiles=tiles,
2460
+ payload_len=declared_payload_len,
2461
+ logical_root=logical_root,
2462
+ )
2463
+
2464
+
2465
+ def _tile_start(counts: "list[int]", index: int) -> "list[int]":
2466
+ coords = [0] * len(counts)
2467
+ remainder = index
2468
+ for axis in range(len(counts) - 1, -1, -1):
2469
+ coords[axis] = remainder % counts[axis]
2470
+ remainder //= counts[axis]
2471
+ return coords
2472
+
2473
+
2474
+ def _decode_tile_planes(
2475
+ stored: bytes, element_size: int, element_len: int
2476
+ ) -> bytes:
2477
+ if element_size == 0 or element_len % element_size != 0:
2478
+ raise Malformed(
2479
+ "tensor tile element length is not a whole number of elements"
2480
+ )
2481
+ elements = element_len // element_size
2482
+ planes = []
2483
+ cursor = 0
2484
+ for _ in range(element_size):
2485
+ if cursor > len(stored):
2486
+ raise Corrupt("tensor tile plane is truncated")
2487
+ plane, consumed = huffman.decode_prefix(stored[cursor:], elements)
2488
+ cursor += consumed
2489
+ planes.append(plane)
2490
+ if cursor != len(stored):
2491
+ raise Corrupt(
2492
+ "tensor tile planes leave trailing bytes inside the tile range"
2493
+ )
2494
+ out = bytearray(element_len)
2495
+ for plane_index, plane in enumerate(planes):
2496
+ for element, byte in enumerate(plane):
2497
+ out[element * element_size + plane_index] = byte
2498
+ return bytes(out)
2499
+
2500
+
2501
+ def _iter_coords(extent: "list[int]"):
2502
+ if not extent:
2503
+ yield []
2504
+ return
2505
+ def walk(dimension: int, prefix: "list[int]"):
2506
+ if dimension == len(extent):
2507
+ yield list(prefix)
2508
+ return
2509
+ for position in range(extent[dimension]):
2510
+ prefix.append(position)
2511
+ yield from walk(dimension + 1, prefix)
2512
+ prefix.pop()
2513
+
2514
+ yield from walk(0, [])
2515
+
2516
+
2517
+ def _flat_index(coords: "list[int]", shape: "list[int]") -> int:
2518
+ index = 0
2519
+ for coord, dim in zip(coords, shape):
2520
+ index = index * dim + coord
2521
+ return index
2522
+
2523
+
2524
+ def decode_tensor_tiled(frame, limits: Limits) -> bytes:
2525
+ descriptor_bytes, payload = _split_prefix(frame.body)
2526
+ meta = parse_tiled_descriptor(
2527
+ descriptor_bytes, frame.header.decoded_len, len(payload), limits
2528
+ )
2529
+ element_size = TENSOR_DTYPE_SIZES[meta.dtype]
2530
+ total = meta.element_count * element_size
2531
+ out = bytearray(total)
2532
+ tiles_bytes = 0
2533
+ for entry in meta.tiles:
2534
+ start = entry.offset
2535
+ end = start + entry.compressed_len
2536
+ if end > len(payload):
2537
+ raise Corrupt("tensor tile range is outside the payload")
2538
+ stored = payload[start:end]
2539
+ if len(stored) != entry.compressed_len:
2540
+ raise Corrupt("tensor tile stored range length is inconsistent")
2541
+ tile = _decode_tile_planes(stored, element_size, entry.element_len)
2542
+ root = tile_root(meta.dtype, entry.dims, tile, limits)
2543
+ if root != entry.tile_root:
2544
+ raise Corrupt(
2545
+ "tensor tile %d root does not match the descriptor" % entry.index
2546
+ )
2547
+ tiles_bytes += len(tile)
2548
+ tile_coords = _tile_start(meta.tile_counts, entry.index)
2549
+ dst_start = [
2550
+ coord * meta.tile_shape[axis]
2551
+ for axis, coord in enumerate(tile_coords)
2552
+ ]
2553
+ for local in _iter_coords(entry.dims):
2554
+ src = _flat_index(local, entry.dims) * element_size
2555
+ full = [
2556
+ dst_start[axis] + local[axis]
2557
+ for axis in range(len(meta.shape))
2558
+ ]
2559
+ dst = _flat_index(full, meta.shape) * element_size
2560
+ out[dst : dst + element_size] = tile[src : src + element_size]
2561
+ if tiles_bytes != total:
2562
+ raise Corrupt(
2563
+ "tiled tensor decoded to %d bytes but the descriptor requires %d"
2564
+ % (tiles_bytes, total)
2565
+ )
2566
+ logical = tensor_logical_root(
2567
+ meta.dtype, meta.shape, meta.element_count, bytes(out), limits
2568
+ )
2569
+ if logical != meta.logical_root:
2570
+ raise Corrupt(
2571
+ "tiled tensor logical root does not match the descriptor"
2572
+ )
2573
+ return bytes(out)
2574
+
2575
+
2576
+ # ---------------------------------------------------------------------------
2577
+ # Mode 17 -- FLOAT_TILED
2578
+ # ---------------------------------------------------------------------------
2579
+
2580
+ FLOAT_TILED_SCHEMA = "weft.ai.float_tiled.v1"
2581
+ FLOAT_TILED_DTYPES = ("bf16", "fp16", "fp32")
2582
+ # The descriptor uses the AI-codec spelling; the WEFT-TENSOR identity uses `f16`.
2583
+ FLOAT_TILED_LOGICAL_DTYPE = {"bf16": "bf16", "fp16": "f16", "fp32": "f32"}
2584
+
2585
+ _FLOAT_TILED_KEYS = (
2586
+ "schema",
2587
+ "dtype",
2588
+ "shape",
2589
+ "element_count",
2590
+ "layout",
2591
+ "tile_shape",
2592
+ "tile_counts",
2593
+ "exponent_coder",
2594
+ "tiles",
2595
+ "payload_len",
2596
+ "logical_root",
2597
+ )
2598
+
2599
+ _FLOAT_TILE_KEYS = (
2600
+ "index",
2601
+ "dims",
2602
+ "offset",
2603
+ "compressed_len",
2604
+ "exponent_plane_len",
2605
+ "mantissa_plane_len",
2606
+ "element_len",
2607
+ "tile_root",
2608
+ )
2609
+
2610
+
2611
+ @dataclass
2612
+ class FloatTileEntry:
2613
+ index: int
2614
+ dims: "list[int]"
2615
+ offset: int
2616
+ compressed_len: int
2617
+ exponent_plane_len: int
2618
+ mantissa_plane_len: int
2619
+ element_len: int
2620
+ tile_root: bytes
2621
+
2622
+
2623
+ @dataclass
2624
+ class FloatTiledMeta:
2625
+ dtype: str
2626
+ shape: "list[int]"
2627
+ element_count: int
2628
+ tile_shape: "list[int]"
2629
+ tile_counts: "list[int]"
2630
+ tiles: "list[FloatTileEntry]"
2631
+ payload_len: int
2632
+ logical_root: bytes
2633
+
2634
+
2635
+ def float_tile_root(
2636
+ dtype: str,
2637
+ dims: "list[int]",
2638
+ element_len: int,
2639
+ exponent_plane_len: int,
2640
+ mantissa_plane_len: int,
2641
+ tile_bytes: bytes,
2642
+ limits: Limits,
2643
+ ) -> bytes:
2644
+ record = {
2645
+ "dtype": dtype,
2646
+ "dims": _u64_vec(dims),
2647
+ "element_len": wcb.U64(element_len),
2648
+ "exponent_plane_len": wcb.U64(exponent_plane_len),
2649
+ "mantissa_plane_len": wcb.U64(mantissa_plane_len),
2650
+ "values_hash": values_hash(tile_bytes),
2651
+ }
2652
+ return hash_domain(DOMAIN_TENSOR_FLOAT_TILE, wcb.encode(record, limits))
2653
+
2654
+
2655
+ def parse_float_tiled_descriptor(
2656
+ descriptor_bytes: bytes, decoded_len: int, payload_len: int, limits: Limits
2657
+ ) -> FloatTiledMeta:
2658
+ map_value = _map(_decode_descriptor(descriptor_bytes, limits), "float tiled tensor")
2659
+ _exact_keys(map_value, _FLOAT_TILED_KEYS, "float tiled tensor")
2660
+ schema = _text(map_value, "schema")
2661
+ if schema != FLOAT_TILED_SCHEMA:
2662
+ raise Malformed("unsupported float tiled tensor schema `%s`" % schema)
2663
+ dtype = _text(map_value, "dtype")
2664
+ if dtype not in FLOAT_FIELD_BITS:
2665
+ raise Malformed(
2666
+ "unknown float tiled tensor dtype `%s` (v1 allows bf16, fp16, "
2667
+ "fp32, fp8-e4m3, fp8-e5m2)" % dtype
2668
+ )
2669
+ if dtype not in FLOAT_TILED_DTYPES:
2670
+ raise UnsupportedFeature(
2671
+ "float tiled tensor dtype `%s` is not implemented in v1" % dtype
2672
+ )
2673
+ layout = _text(map_value, "layout")
2674
+ if layout != TILE_LAYOUT:
2675
+ raise Malformed(
2676
+ "float tiled tensor layout `%s` is specified but not implemented "
2677
+ "in v1 (only %s)" % (layout, TILE_LAYOUT)
2678
+ )
2679
+ coder = _text(map_value, "exponent_coder")
2680
+ if coder != PLANE_CODEC:
2681
+ raise Malformed(
2682
+ "unsupported float tiled tensor exponent coder `%s` (v1 allows `%s`)"
2683
+ % (coder, PLANE_CODEC)
2684
+ )
2685
+ shape = _shape(map_value)
2686
+ element_count = _validate_shape(shape, limits, "float tiled tensor")
2687
+ declared_count = _u64(map_value, "element_count")
2688
+ if declared_count != element_count:
2689
+ raise Malformed(
2690
+ "float tiled tensor element_count %d does not match shape product %d"
2691
+ % (declared_count, element_count)
2692
+ )
2693
+ sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[dtype]
2694
+ element_bits = sign_bits + exponent_bits + mantissa_bits
2695
+ if element_bits % 8 != 0:
2696
+ raise Malformed("float tiled tensor element width is not byte-aligned")
2697
+ element_size = element_bits // 8
2698
+ decoded_bytes = _checked_mul(
2699
+ element_count, element_size, "float tiled tensor decoded bytes"
2700
+ )
2701
+ _check_len(decoded_bytes, limits.max_decoded_bytes, "float tiled tensor decoded bytes")
2702
+ if decoded_len != decoded_bytes:
2703
+ raise Corrupt(
2704
+ "float tiled tensor decoded_len is %d but the descriptor requires %d"
2705
+ % (decoded_len, decoded_bytes)
2706
+ )
2707
+ tile_shape = _array_u64(map_value, "tile_shape")
2708
+ counts = _tile_counts(shape, tile_shape)
2709
+ declared_counts = _array_u64(map_value, "tile_counts")
2710
+ if declared_counts != counts:
2711
+ raise Malformed(
2712
+ "float tiled tensor tile_counts %r do not match shape/tile_shape %r"
2713
+ % (declared_counts, counts)
2714
+ )
2715
+ tile_elements = _validate_shape(tile_shape, limits, "float tensor tile")
2716
+ tile_element_len = _checked_mul(
2717
+ tile_elements, element_size, "float tensor tile byte length"
2718
+ )
2719
+ tile_count = 1
2720
+ for count in counts:
2721
+ tile_count = _checked_mul(tile_count, count, "float tensor tile count")
2722
+ if tile_count > limits.max_structured_items:
2723
+ raise LimitExceeded(
2724
+ "float tensor tile count %d exceeds limit %d"
2725
+ % (tile_count, limits.max_structured_items)
2726
+ )
2727
+ declared_payload_len = _u64(map_value, "payload_len")
2728
+ if declared_payload_len != payload_len:
2729
+ # The frame body is the authority; a declared length that disagrees is
2730
+ # a truncated or over-long payload, which is corruption.
2731
+ raise Corrupt(
2732
+ "float tiled tensor payload is %d bytes but metadata declares %d"
2733
+ % (payload_len, declared_payload_len)
2734
+ )
2735
+ mantissa_total_bits = _checked_mul(
2736
+ tile_elements, sign_bits + mantissa_bits, "float tensor tile mantissa bits"
2737
+ )
2738
+ expected_tile_mantissa = _checked_add(
2739
+ mantissa_total_bits, 7, "float tensor tile mantissa length"
2740
+ ) // 8
2741
+ entries = map_value.get("tiles")
2742
+ if not isinstance(entries, list):
2743
+ raise Malformed("float tiled tensor `tiles` must be an array")
2744
+ if len(entries) != tile_count:
2745
+ raise Malformed(
2746
+ "float tiled tensor has %d tile entries but the grid requires %d"
2747
+ % (len(entries), tile_count)
2748
+ )
2749
+ tiles = []
2750
+ expected_offset = 0
2751
+ for position, item in enumerate(entries):
2752
+ entry = _map(item, "float tensor tile")
2753
+ _exact_keys(entry, _FLOAT_TILE_KEYS, "float tensor tile")
2754
+ index = _u64(entry, "index")
2755
+ if index != position:
2756
+ raise Malformed(
2757
+ "float tiled tensor entry %d declares index %d; entries must be "
2758
+ "in canonical order" % (position, index)
2759
+ )
2760
+ if index >= tile_count:
2761
+ raise Malformed(
2762
+ "float tiled tensor tile index %d is outside the %d-tile grid"
2763
+ % (index, tile_count)
2764
+ )
2765
+ dims = _array_u64(entry, "dims")
2766
+ if dims != tile_shape:
2767
+ raise Malformed(
2768
+ "float tensor tile %d dims %r do not equal tile_shape %r; "
2769
+ "partial tiles are not defined in v1" % (position, dims, tile_shape)
2770
+ )
2771
+ declared_element_len = _u64(entry, "element_len")
2772
+ if declared_element_len != tile_element_len:
2773
+ raise Malformed(
2774
+ "float tensor tile %d element_len %d does not match %d"
2775
+ % (position, declared_element_len, tile_element_len)
2776
+ )
2777
+ offset = _u64(entry, "offset")
2778
+ if offset != expected_offset:
2779
+ raise Malformed(
2780
+ "float tensor tile %d offset %d is not the contiguous offset %d"
2781
+ % (position, offset, expected_offset)
2782
+ )
2783
+ compressed_len = _u64(entry, "compressed_len")
2784
+ if compressed_len == 0:
2785
+ raise Malformed("float tensor tile %d has an empty stored range" % position)
2786
+ exponent_plane_len = _u64(entry, "exponent_plane_len")
2787
+ if exponent_plane_len < 256:
2788
+ raise Malformed(
2789
+ "float tensor tile %d exponent_plane_len %d is below the "
2790
+ "Huffman table minimum" % (position, exponent_plane_len)
2791
+ )
2792
+ mantissa_plane_len = _u64(entry, "mantissa_plane_len")
2793
+ if mantissa_plane_len != expected_tile_mantissa:
2794
+ raise Malformed(
2795
+ "float tensor tile %d mantissa_plane_len %d does not match the "
2796
+ "tile arithmetic %d"
2797
+ % (position, mantissa_plane_len, expected_tile_mantissa)
2798
+ )
2799
+ if exponent_plane_len + mantissa_plane_len != compressed_len:
2800
+ raise Malformed(
2801
+ "float tensor tile %d plane lengths %d + %d do not sum to "
2802
+ "compressed_len %d"
2803
+ % (
2804
+ position,
2805
+ exponent_plane_len,
2806
+ mantissa_plane_len,
2807
+ compressed_len,
2808
+ )
2809
+ )
2810
+ expected_offset = _checked_add(
2811
+ expected_offset, compressed_len, "float tensor tile offset"
2812
+ )
2813
+ root = _bytes_of(entry, "tile_root")
2814
+ if len(root) != HASH_LEN:
2815
+ raise Malformed("float tensor tile_root must be 32 bytes")
2816
+ tiles.append(
2817
+ FloatTileEntry(
2818
+ index=index,
2819
+ dims=dims,
2820
+ offset=offset,
2821
+ compressed_len=compressed_len,
2822
+ exponent_plane_len=exponent_plane_len,
2823
+ mantissa_plane_len=mantissa_plane_len,
2824
+ element_len=declared_element_len,
2825
+ tile_root=root,
2826
+ )
2827
+ )
2828
+ if expected_offset != declared_payload_len:
2829
+ raise Corrupt(
2830
+ "float tensor tile ranges cover %d bytes but payload_len is %d"
2831
+ % (expected_offset, declared_payload_len)
2832
+ )
2833
+ logical_root = _bytes_of(map_value, "logical_root")
2834
+ if len(logical_root) != HASH_LEN:
2835
+ raise Malformed("float tiled tensor logical_root must be 32 bytes")
2836
+ return FloatTiledMeta(
2837
+ dtype=dtype,
2838
+ shape=shape,
2839
+ element_count=element_count,
2840
+ tile_shape=tile_shape,
2841
+ tile_counts=counts,
2842
+ tiles=tiles,
2843
+ payload_len=declared_payload_len,
2844
+ logical_root=logical_root,
2845
+ )
2846
+
2847
+
2848
+ def decode_float_tiled(frame, limits: Limits) -> bytes:
2849
+ descriptor_bytes, payload = _split_prefix(frame.body)
2850
+ meta = parse_float_tiled_descriptor(
2851
+ descriptor_bytes, frame.header.decoded_len, len(payload), limits
2852
+ )
2853
+ sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[meta.dtype]
2854
+ element_size = (sign_bits + exponent_bits + mantissa_bits) // 8
2855
+ total = meta.element_count * element_size
2856
+ out = bytearray(total)
2857
+ tiles_bytes = 0
2858
+ for entry in meta.tiles:
2859
+ start = entry.offset
2860
+ end = start + entry.compressed_len
2861
+ if end > len(payload):
2862
+ raise Corrupt("float tensor tile range is outside the payload")
2863
+ stored = payload[start:end]
2864
+ exponent_compressed = stored[: entry.exponent_plane_len]
2865
+ mantissa_plane = stored[entry.exponent_plane_len :]
2866
+ elements = entry.element_len // element_size
2867
+ exponent_plane = huffman.decode(exponent_compressed, elements)
2868
+ tile = _join_float_planes(
2869
+ exponent_plane,
2870
+ mantissa_plane,
2871
+ meta.dtype,
2872
+ elements,
2873
+ limits,
2874
+ )
2875
+ root = float_tile_root(
2876
+ meta.dtype,
2877
+ entry.dims,
2878
+ entry.element_len,
2879
+ entry.exponent_plane_len,
2880
+ entry.mantissa_plane_len,
2881
+ tile,
2882
+ limits,
2883
+ )
2884
+ if root != entry.tile_root:
2885
+ raise Corrupt(
2886
+ "float tensor tile %d root does not match the descriptor"
2887
+ % entry.index
2888
+ )
2889
+ tiles_bytes += len(tile)
2890
+ tile_coords = _tile_start(meta.tile_counts, entry.index)
2891
+ dst_start = [
2892
+ coord * meta.tile_shape[axis] for axis, coord in enumerate(tile_coords)
2893
+ ]
2894
+ for local in _iter_coords(entry.dims):
2895
+ src = _flat_index(local, entry.dims) * element_size
2896
+ full = [dst_start[axis] + local[axis] for axis in range(len(meta.shape))]
2897
+ dst = _flat_index(full, meta.shape) * element_size
2898
+ out[dst : dst + element_size] = tile[src : src + element_size]
2899
+ if tiles_bytes != total:
2900
+ raise Corrupt(
2901
+ "float tiled tensor decoded to %d bytes but the descriptor "
2902
+ "requires %d" % (tiles_bytes, total)
2903
+ )
2904
+ logical = tensor_logical_root(
2905
+ FLOAT_TILED_LOGICAL_DTYPE[meta.dtype],
2906
+ meta.shape,
2907
+ meta.element_count,
2908
+ bytes(out),
2909
+ limits,
2910
+ )
2911
+ if logical != meta.logical_root:
2912
+ raise Corrupt(
2913
+ "float tiled tensor logical root does not match the descriptor"
2914
+ )
2915
+ return bytes(out)
2916
+
2917
+
2918
+ # ---------------------------------------------------------------------------
2919
+ # Mode 18 -- FLOAT_TILED_STREAMS
2920
+ # ---------------------------------------------------------------------------
2921
+
2922
+ FLOAT_TILED_STREAMS_SCHEMA = "weft.ai.tensor_tiled_streams.v1"
2923
+ FLOAT_TILED_STREAMS_MIN = 1
2924
+ FLOAT_TILED_STREAMS_MAX = 16
2925
+
2926
+ _FLOAT_TILED_STREAMS_KEYS = (
2927
+ "schema",
2928
+ "dtype",
2929
+ "shape",
2930
+ "element_count",
2931
+ "layout",
2932
+ "tile_shape",
2933
+ "tile_counts",
2934
+ "exponent_coder",
2935
+ "n_streams",
2936
+ "tiles",
2937
+ "payload_len",
2938
+ "logical_root",
2939
+ )
2940
+
2941
+ _FLOAT_TILE_STREAMS_KEYS = (
2942
+ "index",
2943
+ "dims",
2944
+ "offset",
2945
+ "compressed_len",
2946
+ "exponent_plane_len",
2947
+ "mantissa_plane_len",
2948
+ "element_len",
2949
+ "stream_offsets",
2950
+ "tile_root",
2951
+ )
2952
+
2953
+
2954
+ @dataclass
2955
+ class FloatTileStreamsEntry:
2956
+ index: int
2957
+ dims: "list[int]"
2958
+ offset: int
2959
+ compressed_len: int
2960
+ exponent_plane_len: int
2961
+ mantissa_plane_len: int
2962
+ element_len: int
2963
+ stream_offsets: "list[int]"
2964
+ tile_root: bytes
2965
+
2966
+
2967
+ @dataclass
2968
+ class FloatTiledStreamsMeta:
2969
+ dtype: str
2970
+ shape: "list[int]"
2971
+ element_count: int
2972
+ tile_shape: "list[int]"
2973
+ tile_counts: "list[int]"
2974
+ n_streams: int
2975
+ tiles: "list[FloatTileStreamsEntry]"
2976
+ payload_len: int
2977
+ logical_root: bytes
2978
+
2979
+
2980
+ def interleaved_stream_len(total: int, stream: int, n_streams: int) -> int:
2981
+ if n_streams <= 0 or total <= stream:
2982
+ return 0
2983
+ return (total - stream + n_streams - 1) // n_streams
2984
+
2985
+
2986
+ def float_tile_streams_root(
2987
+ dtype: str,
2988
+ dims: "list[int]",
2989
+ element_len: int,
2990
+ exponent_plane_len: int,
2991
+ mantissa_plane_len: int,
2992
+ n_streams: int,
2993
+ stream_offsets: "list[int]",
2994
+ tile_bytes: bytes,
2995
+ limits: Limits,
2996
+ ) -> bytes:
2997
+ record = {
2998
+ "dtype": dtype,
2999
+ "dims": _u64_vec(dims),
3000
+ "element_len": wcb.U64(element_len),
3001
+ "exponent_plane_len": wcb.U64(exponent_plane_len),
3002
+ "mantissa_plane_len": wcb.U64(mantissa_plane_len),
3003
+ "n_streams": wcb.U64(n_streams),
3004
+ "stream_offsets": _u64_vec(stream_offsets),
3005
+ "values_hash": values_hash(tile_bytes),
3006
+ }
3007
+ return hash_domain(DOMAIN_TENSOR_FLOAT_TILE_STREAMS, wcb.encode(record, limits))
3008
+
3009
+
3010
+ def parse_float_tiled_streams_descriptor(
3011
+ descriptor_bytes: bytes, decoded_len: int, payload_len: int, limits: Limits
3012
+ ) -> FloatTiledStreamsMeta:
3013
+ map_value = _map(
3014
+ _decode_descriptor(descriptor_bytes, limits), "float tiled streams tensor"
3015
+ )
3016
+ _exact_keys(map_value, _FLOAT_TILED_STREAMS_KEYS, "float tiled streams tensor")
3017
+ schema = _text(map_value, "schema")
3018
+ if schema != FLOAT_TILED_STREAMS_SCHEMA:
3019
+ raise Malformed("unsupported float tiled streams schema `%s`" % schema)
3020
+ dtype = _text(map_value, "dtype")
3021
+ if dtype not in FLOAT_FIELD_BITS:
3022
+ raise Malformed("unknown float tiled streams dtype `%s`" % dtype)
3023
+ if dtype not in FLOAT_TILED_DTYPES:
3024
+ raise UnsupportedFeature(
3025
+ "float tiled streams dtype `%s` is not implemented in v1" % dtype
3026
+ )
3027
+ layout = _text(map_value, "layout")
3028
+ if layout != TILE_LAYOUT:
3029
+ raise Malformed(
3030
+ "float tiled streams layout `%s` is specified but not implemented "
3031
+ "in v1 (only %s)" % (layout, TILE_LAYOUT)
3032
+ )
3033
+ coder = _text(map_value, "exponent_coder")
3034
+ if coder != PLANE_CODEC:
3035
+ raise Malformed(
3036
+ "unsupported float tiled streams exponent coder `%s` (v1 allows `%s`)"
3037
+ % (coder, PLANE_CODEC)
3038
+ )
3039
+ n_streams = _u64(map_value, "n_streams")
3040
+ if not FLOAT_TILED_STREAMS_MIN <= n_streams <= FLOAT_TILED_STREAMS_MAX:
3041
+ raise Malformed(
3042
+ "n_streams %d is outside %d..=%d"
3043
+ % (n_streams, FLOAT_TILED_STREAMS_MIN, FLOAT_TILED_STREAMS_MAX)
3044
+ )
3045
+ shape = _shape(map_value)
3046
+ element_count = _validate_shape(shape, limits, "float tiled streams tensor")
3047
+ declared_count = _u64(map_value, "element_count")
3048
+ if declared_count != element_count:
3049
+ raise Malformed(
3050
+ "float tiled streams element_count %d does not match shape product %d"
3051
+ % (declared_count, element_count)
3052
+ )
3053
+ sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[dtype]
3054
+ element_size = (sign_bits + exponent_bits + mantissa_bits) // 8
3055
+ decoded_bytes = _checked_mul(
3056
+ element_count, element_size, "float tiled streams decoded bytes"
3057
+ )
3058
+ _check_len(
3059
+ decoded_bytes, limits.max_decoded_bytes, "float tiled streams decoded bytes"
3060
+ )
3061
+ if decoded_len != decoded_bytes:
3062
+ raise Corrupt(
3063
+ "float tiled streams decoded_len is %d but the descriptor requires %d"
3064
+ % (decoded_len, decoded_bytes)
3065
+ )
3066
+ tile_shape = _array_u64(map_value, "tile_shape")
3067
+ counts = _tile_counts(shape, tile_shape)
3068
+ if _array_u64(map_value, "tile_counts") != counts:
3069
+ raise Malformed(
3070
+ "float tiled streams tile_counts %r do not match shape/tile_shape %r"
3071
+ % (_array_u64(map_value, "tile_counts"), counts)
3072
+ )
3073
+ tile_elements = _validate_shape(tile_shape, limits, "float tensor tile")
3074
+ tile_element_len = _checked_mul(
3075
+ tile_elements, element_size, "float tensor tile byte length"
3076
+ )
3077
+ tile_count = 1
3078
+ for count in counts:
3079
+ tile_count = _checked_mul(tile_count, count, "float tensor tile count")
3080
+ if tile_count > limits.max_structured_items:
3081
+ raise LimitExceeded(
3082
+ "float tensor tile count %d exceeds limit %d"
3083
+ % (tile_count, limits.max_structured_items)
3084
+ )
3085
+ mantissa_total_bits = _checked_mul(
3086
+ tile_elements, sign_bits + mantissa_bits, "float tensor tile mantissa bits"
3087
+ )
3088
+ expected_tile_mantissa = _checked_add(
3089
+ mantissa_total_bits, 7, "float tensor tile mantissa length"
3090
+ ) // 8
3091
+ declared_payload_len = _u64(map_value, "payload_len")
3092
+ if declared_payload_len != payload_len:
3093
+ raise Corrupt(
3094
+ "float tiled streams payload is %d bytes but metadata declares %d"
3095
+ % (payload_len, declared_payload_len)
3096
+ )
3097
+ entries = map_value.get("tiles")
3098
+ if not isinstance(entries, list):
3099
+ raise Malformed("float tiled streams `tiles` must be an array")
3100
+ if len(entries) != tile_count:
3101
+ raise Malformed(
3102
+ "float tiled streams has %d tile entries but the grid requires %d"
3103
+ % (len(entries), tile_count)
3104
+ )
3105
+ tiles = []
3106
+ expected_offset = 0
3107
+ for position, item in enumerate(entries):
3108
+ entry = _map(item, "float tensor streams tile")
3109
+ _exact_keys(entry, _FLOAT_TILE_STREAMS_KEYS, "float tensor streams tile")
3110
+ index = _u64(entry, "index")
3111
+ if index != position:
3112
+ raise Malformed(
3113
+ "float tiled streams entry %d declares index %d; entries must be "
3114
+ "in canonical order" % (position, index)
3115
+ )
3116
+ dims = _array_u64(entry, "dims")
3117
+ if dims != tile_shape:
3118
+ raise Malformed(
3119
+ "float tensor tile %d dims %r do not equal tile_shape %r; "
3120
+ "partial tiles are not defined in v1" % (position, dims, tile_shape)
3121
+ )
3122
+ declared_element_len = _u64(entry, "element_len")
3123
+ if declared_element_len != tile_element_len:
3124
+ raise Malformed(
3125
+ "float tensor tile %d element_len %d does not match %d"
3126
+ % (position, declared_element_len, tile_element_len)
3127
+ )
3128
+ offset = _u64(entry, "offset")
3129
+ if offset != expected_offset:
3130
+ raise Malformed(
3131
+ "float tensor tile %d offset %d is not the contiguous offset %d"
3132
+ % (position, offset, expected_offset)
3133
+ )
3134
+ compressed_len = _u64(entry, "compressed_len")
3135
+ exponent_plane_len = _u64(entry, "exponent_plane_len")
3136
+ mantissa_plane_len = _u64(entry, "mantissa_plane_len")
3137
+ if mantissa_plane_len != expected_tile_mantissa:
3138
+ raise Malformed(
3139
+ "float tensor tile %d mantissa_plane_len %d does not match the "
3140
+ "tile arithmetic %d"
3141
+ % (position, mantissa_plane_len, expected_tile_mantissa)
3142
+ )
3143
+ if exponent_plane_len + mantissa_plane_len != compressed_len:
3144
+ raise Malformed(
3145
+ "float tensor tile %d plane lengths do not sum to compressed_len"
3146
+ % position
3147
+ )
3148
+ expected_offset = _checked_add(
3149
+ expected_offset, compressed_len, "float tensor tile offset"
3150
+ )
3151
+ stream_offsets = _array_u64(entry, "stream_offsets")
3152
+ if len(stream_offsets) != 2 * n_streams:
3153
+ raise Malformed(
3154
+ "float tile %d has %d stream offsets but n_streams %d requires %d"
3155
+ % (position, len(stream_offsets), n_streams, 2 * n_streams)
3156
+ )
3157
+ if stream_offsets[0] != 0:
3158
+ raise Malformed(
3159
+ "float tile %d stream_offsets[0] is %d not 0"
3160
+ % (position, stream_offsets[0])
3161
+ )
3162
+ for earlier, later in zip(stream_offsets, stream_offsets[1:]):
3163
+ if later < earlier:
3164
+ raise Malformed(
3165
+ "float tile %d stream_offsets are not monotonic: %d then %d"
3166
+ % (position, earlier, later)
3167
+ )
3168
+ if stream_offsets[-1] > compressed_len:
3169
+ raise Malformed(
3170
+ "float tile %d stream offset %d is past the %d-byte tile payload"
3171
+ % (position, stream_offsets[-1], compressed_len)
3172
+ )
3173
+ if stream_offsets[n_streams] != exponent_plane_len:
3174
+ raise Malformed(
3175
+ "float tile %d stream boundary %d does not equal exponent_plane_len %d"
3176
+ % (position, stream_offsets[n_streams], exponent_plane_len)
3177
+ )
3178
+ for stream in range(n_streams):
3179
+ end = (
3180
+ stream_offsets[stream + 1]
3181
+ if stream + 1 < n_streams
3182
+ else exponent_plane_len
3183
+ )
3184
+ if end - stream_offsets[stream] < 256:
3185
+ raise Malformed(
3186
+ "float tile %d exponent stream %d is shorter than its "
3187
+ "256-byte code-length table" % (position, stream)
3188
+ )
3189
+ for stream in range(n_streams):
3190
+ start = stream_offsets[n_streams + stream] - exponent_plane_len
3191
+ end = (
3192
+ stream_offsets[n_streams + stream + 1] - exponent_plane_len
3193
+ if stream + 1 < n_streams
3194
+ else mantissa_plane_len
3195
+ )
3196
+ expected = interleaved_stream_len(
3197
+ mantissa_plane_len, stream, n_streams
3198
+ )
3199
+ if end - start != expected:
3200
+ raise Malformed(
3201
+ "float tile %d mantissa stream %d is %d bytes but the "
3202
+ "interleave requires %d"
3203
+ % (position, stream, end - start, expected)
3204
+ )
3205
+ root = _bytes_of(entry, "tile_root")
3206
+ if len(root) != HASH_LEN:
3207
+ raise Malformed("float tensor tile_root must be 32 bytes")
3208
+ tiles.append(
3209
+ FloatTileStreamsEntry(
3210
+ index=index,
3211
+ dims=dims,
3212
+ offset=offset,
3213
+ compressed_len=compressed_len,
3214
+ exponent_plane_len=exponent_plane_len,
3215
+ mantissa_plane_len=mantissa_plane_len,
3216
+ element_len=declared_element_len,
3217
+ stream_offsets=stream_offsets,
3218
+ tile_root=root,
3219
+ )
3220
+ )
3221
+ if expected_offset != declared_payload_len:
3222
+ raise Corrupt(
3223
+ "float tensor tile ranges cover %d bytes but payload_len is %d"
3224
+ % (expected_offset, declared_payload_len)
3225
+ )
3226
+ logical_root = _bytes_of(map_value, "logical_root")
3227
+ if len(logical_root) != HASH_LEN:
3228
+ raise Malformed("float tiled streams logical_root must be 32 bytes")
3229
+ return FloatTiledStreamsMeta(
3230
+ dtype=dtype,
3231
+ shape=shape,
3232
+ element_count=element_count,
3233
+ tile_shape=tile_shape,
3234
+ tile_counts=counts,
3235
+ n_streams=n_streams,
3236
+ tiles=tiles,
3237
+ payload_len=declared_payload_len,
3238
+ logical_root=logical_root,
3239
+ )
3240
+
3241
+
3242
+ def _deinterleave(streams: "list[bytes]", total: int, what: str) -> bytes:
3243
+ n_streams = len(streams)
3244
+ if n_streams == 0:
3245
+ raise Malformed("%s declares zero streams" % what)
3246
+ out = bytearray(total)
3247
+ for index in range(total):
3248
+ stream = streams[index % n_streams]
3249
+ position = index // n_streams
3250
+ if position >= len(stream):
3251
+ raise Corrupt("%s stream ended before its interleaved plane" % what)
3252
+ out[index] = stream[position]
3253
+ return bytes(out)
3254
+
3255
+
3256
+ def decode_float_tiled_streams(frame, limits: Limits) -> bytes:
3257
+ descriptor_bytes, payload = _split_prefix(frame.body)
3258
+ meta = parse_float_tiled_streams_descriptor(
3259
+ descriptor_bytes, frame.header.decoded_len, len(payload), limits
3260
+ )
3261
+ sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[meta.dtype]
3262
+ element_size = (sign_bits + exponent_bits + mantissa_bits) // 8
3263
+ total = meta.element_count * element_size
3264
+ out = bytearray(total)
3265
+ tiles_bytes = 0
3266
+ n_streams = meta.n_streams
3267
+ for entry in meta.tiles:
3268
+ start = entry.offset
3269
+ end = start + entry.compressed_len
3270
+ if end > len(payload):
3271
+ raise Corrupt("float tensor tile range is outside the payload")
3272
+ stored = payload[start:end]
3273
+ if entry.exponent_plane_len > len(stored):
3274
+ raise Corrupt("float tile exponent plane is truncated")
3275
+ exponent_stored = stored[: entry.exponent_plane_len]
3276
+ mantissa_stored = stored[entry.exponent_plane_len :]
3277
+ elements = entry.element_len // element_size
3278
+ exponent_parts = []
3279
+ for stream in range(n_streams):
3280
+ s_start = entry.stream_offsets[stream]
3281
+ s_end = (
3282
+ entry.stream_offsets[stream + 1]
3283
+ if stream + 1 < n_streams
3284
+ else entry.exponent_plane_len
3285
+ )
3286
+ if s_end > len(exponent_stored) or s_start > s_end:
3287
+ raise Corrupt("float tile exponent stream is outside the tile")
3288
+ expected = interleaved_stream_len(elements, stream, n_streams)
3289
+ exponent_parts.append(
3290
+ huffman.decode(exponent_stored[s_start:s_end], expected)
3291
+ )
3292
+ exponent_plane = _deinterleave(exponent_parts, elements, "exponent plane")
3293
+ mantissa_parts = []
3294
+ for stream in range(n_streams):
3295
+ s_start = entry.stream_offsets[n_streams + stream] - entry.exponent_plane_len
3296
+ s_end = (
3297
+ entry.stream_offsets[n_streams + stream + 1] - entry.exponent_plane_len
3298
+ if stream + 1 < n_streams
3299
+ else entry.mantissa_plane_len
3300
+ )
3301
+ if s_end > len(mantissa_stored) or s_start > s_end:
3302
+ raise Corrupt("float tile mantissa stream is outside the tile")
3303
+ expected = interleaved_stream_len(
3304
+ entry.mantissa_plane_len, stream, n_streams
3305
+ )
3306
+ if s_end - s_start != expected:
3307
+ raise Corrupt(
3308
+ "float tile %d mantissa stream %d length differs from the "
3309
+ "interleave" % (entry.index, stream)
3310
+ )
3311
+ mantissa_parts.append(mantissa_stored[s_start:s_end])
3312
+ mantissa_plane = _deinterleave(
3313
+ mantissa_parts, entry.mantissa_plane_len, "mantissa plane"
3314
+ )
3315
+ tile = _join_float_planes(
3316
+ exponent_plane, mantissa_plane, meta.dtype, elements, limits
3317
+ )
3318
+ root = float_tile_streams_root(
3319
+ meta.dtype,
3320
+ entry.dims,
3321
+ entry.element_len,
3322
+ entry.exponent_plane_len,
3323
+ entry.mantissa_plane_len,
3324
+ n_streams,
3325
+ entry.stream_offsets,
3326
+ tile,
3327
+ limits,
3328
+ )
3329
+ if root != entry.tile_root:
3330
+ raise Corrupt(
3331
+ "float tensor tile %d stream root does not match the descriptor"
3332
+ % entry.index
3333
+ )
3334
+ tiles_bytes += len(tile)
3335
+ tile_coords = _tile_start(meta.tile_counts, entry.index)
3336
+ dst_start = [
3337
+ coord * meta.tile_shape[axis] for axis, coord in enumerate(tile_coords)
3338
+ ]
3339
+ for local in _iter_coords(entry.dims):
3340
+ src = _flat_index(local, entry.dims) * element_size
3341
+ full = [dst_start[axis] + local[axis] for axis in range(len(meta.shape))]
3342
+ dst = _flat_index(full, meta.shape) * element_size
3343
+ out[dst : dst + element_size] = tile[src : src + element_size]
3344
+ if tiles_bytes != total:
3345
+ raise Corrupt(
3346
+ "float tiled streams decoded to %d bytes but the descriptor requires %d"
3347
+ % (tiles_bytes, total)
3348
+ )
3349
+ logical = tensor_logical_root(
3350
+ FLOAT_TILED_LOGICAL_DTYPE[meta.dtype],
3351
+ meta.shape,
3352
+ meta.element_count,
3353
+ bytes(out),
3354
+ limits,
3355
+ )
3356
+ if logical != meta.logical_root:
3357
+ raise Corrupt(
3358
+ "float tiled streams logical root does not match the descriptor"
3359
+ )
3360
+ return bytes(out)