1use std::collections::BTreeMap;
4
5use serde::{Deserialize, Serialize};
6
7use crate::checkpoint::TensorDtype;
8
9#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12#[non_exhaustive]
13pub enum InputModality {
14 Text,
16 Image,
18 Video,
20 Audio,
22}
23
24impl InputModality {
25 pub const fn as_str(self) -> &'static str {
27 match self {
28 Self::Text => "text",
29 Self::Image => "image",
30 Self::Video => "video",
31 Self::Audio => "audio",
32 }
33 }
34
35 const fn wire_tag(self) -> u32 {
36 match self {
37 Self::Text => 0,
38 Self::Image => 1,
39 Self::Video => 2,
40 Self::Audio => 3,
41 }
42 }
43
44 fn from_wire_tag(tag: u32) -> Result<Self, PreparedInputError> {
45 match tag {
46 0 => Ok(Self::Text),
47 1 => Ok(Self::Image),
48 2 => Ok(Self::Video),
49 3 => Ok(Self::Audio),
50 _ => Err(PreparedInputError::InvalidWireValue {
51 field: "modality",
52 value: tag,
53 }),
54 }
55 }
56}
57
58#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
60#[serde(rename_all = "snake_case")]
61#[non_exhaustive]
62pub enum InputPayloadKind {
63 TokenIds,
65 Tensor,
67 Embeddings,
69}
70
71impl InputPayloadKind {
72 pub const fn accepts(self, modality: InputModality) -> bool {
74 match self {
75 Self::TokenIds => matches!(modality, InputModality::Text),
76 Self::Tensor => !matches!(modality, InputModality::Text),
77 Self::Embeddings => true,
78 }
79 }
80
81 const fn wire_tag(self) -> u32 {
82 match self {
83 Self::TokenIds => 0,
84 Self::Tensor => 1,
85 Self::Embeddings => 2,
86 }
87 }
88
89 fn from_wire_tag(tag: u32) -> Result<Self, PreparedInputError> {
90 match tag {
91 0 => Ok(Self::TokenIds),
92 1 => Ok(Self::Tensor),
93 2 => Ok(Self::Embeddings),
94 _ => Err(PreparedInputError::InvalidWireValue {
95 field: "payload kind",
96 value: tag,
97 }),
98 }
99 }
100}
101
102#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
104#[serde(rename_all = "snake_case")]
105#[non_exhaustive]
106pub enum InputMetadataKey {
107 PatchGrid,
109 PatchPositions,
111 AudioMask,
113}
114
115#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
118#[serde(rename_all = "snake_case")]
119#[non_exhaustive]
120pub enum InputExtent {
121 PatchGrid {
123 time: usize,
125 height: usize,
127 width: usize,
129 },
130 AudioValidFrames(usize),
132}
133
134impl InputExtent {
135 pub const fn accepts(self, modality: InputModality) -> bool {
137 match self {
138 Self::PatchGrid { .. } => {
139 matches!(modality, InputModality::Image | InputModality::Video)
140 }
141 Self::AudioValidFrames(_) => matches!(modality, InputModality::Audio),
142 }
143 }
144
145 const fn wire_tag(self) -> u32 {
146 match self {
147 Self::PatchGrid { .. } => 0,
148 Self::AudioValidFrames(_) => 1,
149 }
150 }
151
152 const fn key(self) -> u32 {
153 self.wire_tag()
154 }
155
156 fn encode_words(self, output: &mut Vec<u32>) -> Result<(), PreparedInputError> {
157 output.push(self.wire_tag());
158 let values: &[usize] = match &self {
159 Self::PatchGrid {
160 time,
161 height,
162 width,
163 } => &[*time, *height, *width],
164 Self::AudioValidFrames(frames) => &[*frames],
165 };
166 for value in values {
167 output.push(
168 u32::try_from(*value)
169 .map_err(|_| PreparedInputError::WireValueOverflow("input extent"))?,
170 );
171 }
172 Ok(())
173 }
174
175 fn decode_words(cursor: &mut WordCursor<'_>) -> Result<Self, PreparedInputError> {
176 match cursor.next("input extent")? {
177 0 => Ok(Self::PatchGrid {
178 time: cursor.usize("patch grid time")?,
179 height: cursor.usize("patch grid height")?,
180 width: cursor.usize("patch grid width")?,
181 }),
182 1 => Ok(Self::AudioValidFrames(cursor.usize("valid audio frames")?)),
183 value => Err(PreparedInputError::InvalidWireValue {
184 field: "input extent",
185 value,
186 }),
187 }
188 }
189}
190
191impl InputMetadataKey {
192 pub const fn accepts(self, modality: InputModality) -> bool {
194 match self {
195 Self::PatchGrid | Self::PatchPositions => {
196 matches!(modality, InputModality::Image | InputModality::Video)
197 }
198 Self::AudioMask => matches!(modality, InputModality::Audio),
199 }
200 }
201
202 const fn wire_tag(self) -> u32 {
203 match self {
204 Self::PatchGrid => 0,
205 Self::PatchPositions => 1,
206 Self::AudioMask => 2,
207 }
208 }
209
210 fn from_wire_tag(tag: u32) -> Result<Self, PreparedInputError> {
211 match tag {
212 0 => Ok(Self::PatchGrid),
213 1 => Ok(Self::PatchPositions),
214 2 => Ok(Self::AudioMask),
215 _ => Err(PreparedInputError::InvalidWireValue {
216 field: "metadata key",
217 value: tag,
218 }),
219 }
220 }
221}
222
223#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
225pub struct InputTensorIdentity {
226 dtype: TensorDtype,
227 shape: Vec<usize>,
228}
229
230impl InputTensorIdentity {
231 pub fn new(dtype: TensorDtype, shape: Vec<usize>) -> Result<Self, PreparedInputError> {
233 if shape.is_empty() || shape.contains(&0) {
234 return Err(PreparedInputError::InvalidTensorShape { shape });
235 }
236 Ok(Self { dtype, shape })
237 }
238
239 pub const fn dtype(&self) -> &TensorDtype {
241 &self.dtype
242 }
243
244 pub fn shape(&self) -> &[usize] {
246 &self.shape
247 }
248
249 fn encode_words(&self, output: &mut Vec<u32>) -> Result<(), PreparedInputError> {
250 encode_dtype(&self.dtype, output)?;
251 output.push(
252 u32::try_from(self.shape.len())
253 .map_err(|_| PreparedInputError::WireValueOverflow("tensor rank"))?,
254 );
255 for dimension in &self.shape {
256 output.push(
257 u32::try_from(*dimension)
258 .map_err(|_| PreparedInputError::WireValueOverflow("tensor dimension"))?,
259 );
260 }
261 Ok(())
262 }
263
264 fn decode_words(cursor: &mut WordCursor<'_>) -> Result<Self, PreparedInputError> {
265 let dtype = decode_dtype(cursor)?;
266 let rank = cursor.usize("tensor rank")?;
267 if rank == 0 || rank > 8 {
268 return Err(PreparedInputError::InvalidWireRank(rank));
269 }
270 let shape = (0..rank)
271 .map(|_| cursor.usize("tensor dimension"))
272 .collect::<Result<Vec<_>, _>>()?;
273 Self::new(dtype, shape)
274 }
275}
276
277#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
279pub struct InputPartDescriptor {
280 modality: InputModality,
281 payload_kind: InputPayloadKind,
282 payload: InputTensorIdentity,
283 metadata: BTreeMap<InputMetadataKey, InputTensorIdentity>,
284 extents: BTreeMap<u32, InputExtent>,
285}
286
287impl InputPartDescriptor {
288 pub fn new(
290 modality: InputModality,
291 payload_kind: InputPayloadKind,
292 payload: InputTensorIdentity,
293 metadata: impl IntoIterator<Item = (InputMetadataKey, InputTensorIdentity)>,
294 ) -> Result<Self, PreparedInputError> {
295 Self::new_with_extents(modality, payload_kind, payload, metadata, [])
296 }
297
298 pub fn new_with_extents(
300 modality: InputModality,
301 payload_kind: InputPayloadKind,
302 payload: InputTensorIdentity,
303 metadata: impl IntoIterator<Item = (InputMetadataKey, InputTensorIdentity)>,
304 extents: impl IntoIterator<Item = InputExtent>,
305 ) -> Result<Self, PreparedInputError> {
306 if !payload_kind.accepts(modality) {
307 return Err(PreparedInputError::IncompatiblePayload {
308 modality,
309 payload: payload_kind,
310 });
311 }
312 let mut typed_metadata = BTreeMap::new();
313 for (key, identity) in metadata {
314 if !key.accepts(modality) {
315 return Err(PreparedInputError::IncompatibleMetadata { modality, key });
316 }
317 if typed_metadata.insert(key, identity).is_some() {
318 return Err(PreparedInputError::DuplicateMetadata { key });
319 }
320 }
321 let mut typed_extents = BTreeMap::new();
322 for extent in extents {
323 if !extent.accepts(modality) {
324 return Err(PreparedInputError::IncompatibleExtent { modality, extent });
325 }
326 if typed_extents.insert(extent.key(), extent).is_some() {
327 return Err(PreparedInputError::DuplicateExtent { extent });
328 }
329 }
330 Ok(Self {
331 modality,
332 payload_kind,
333 payload,
334 metadata: typed_metadata,
335 extents: typed_extents,
336 })
337 }
338
339 pub const fn modality(&self) -> InputModality {
341 self.modality
342 }
343
344 pub const fn payload_kind(&self) -> InputPayloadKind {
346 self.payload_kind
347 }
348
349 pub const fn payload(&self) -> &InputTensorIdentity {
351 &self.payload
352 }
353
354 pub const fn metadata(&self) -> &BTreeMap<InputMetadataKey, InputTensorIdentity> {
356 &self.metadata
357 }
358
359 pub fn extents(&self) -> impl ExactSizeIterator<Item = InputExtent> + '_ {
361 self.extents.values().copied()
362 }
363
364 pub fn metadata_value(&self, key: InputMetadataKey) -> Option<&InputTensorIdentity> {
366 self.metadata.get(&key)
367 }
368
369 pub fn require_metadata(
371 &self,
372 part: usize,
373 key: InputMetadataKey,
374 ) -> Result<&InputTensorIdentity, PreparedInputError> {
375 self.metadata
376 .get(&key)
377 .ok_or(PreparedInputError::MissingMetadata {
378 part,
379 modality: self.modality,
380 key,
381 })
382 }
383}
384
385#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
387pub struct PreparedInputIdentity {
388 parts: Vec<InputPartDescriptor>,
389}
390
391impl PreparedInputIdentity {
392 pub fn new(parts: Vec<InputPartDescriptor>) -> Result<Self, PreparedInputError> {
394 if parts.is_empty() {
395 return Err(PreparedInputError::EmptyInput);
396 }
397 Ok(Self { parts })
398 }
399
400 pub fn parts(&self) -> &[InputPartDescriptor] {
402 &self.parts
403 }
404
405 pub fn logical_metadata_bytes(&self) -> Option<u64> {
409 fn bytes<T>(count: usize) -> Option<u64> {
410 u64::try_from(std::mem::size_of::<T>().checked_mul(count)?).ok()
411 }
412 fn tensor_heap(tensor: &InputTensorIdentity) -> Option<u64> {
413 let shape = bytes::<usize>(tensor.shape.len())?;
414 match &tensor.dtype {
415 TensorDtype::Encoded(name) => shape.checked_add(u64::try_from(name.len()).ok()?),
416 _ => Some(shape),
417 }
418 }
419 let mut total =
420 bytes::<Self>(1)?.checked_add(bytes::<InputPartDescriptor>(self.parts.len())?)?;
421 for part in &self.parts {
422 total = total
423 .checked_add(tensor_heap(&part.payload)?)?
424 .checked_add(bytes::<(InputMetadataKey, InputTensorIdentity)>(
425 part.metadata.len(),
426 )?)?
427 .checked_add(bytes::<(u32, InputExtent)>(part.extents.len())?)?;
428 for tensor in part.metadata.values() {
429 total = total.checked_add(tensor_heap(tensor)?)?;
430 }
431 }
432 Some(total)
433 }
434
435 pub fn len(&self) -> usize {
437 self.parts.len()
438 }
439
440 pub fn is_empty(&self) -> bool {
442 self.parts.is_empty()
443 }
444
445 pub fn encode_words(&self) -> Result<Vec<u32>, PreparedInputError> {
447 let mut output = Vec::new();
448 output.push(
449 u32::try_from(self.parts.len())
450 .map_err(|_| PreparedInputError::WireValueOverflow("part count"))?,
451 );
452 for part in &self.parts {
453 output.extend_from_slice(&[part.modality.wire_tag(), part.payload_kind.wire_tag()]);
454 part.payload.encode_words(&mut output)?;
455 output.push(
456 u32::try_from(part.metadata.len())
457 .map_err(|_| PreparedInputError::WireValueOverflow("metadata count"))?,
458 );
459 for (key, identity) in &part.metadata {
460 output.push(key.wire_tag());
461 identity.encode_words(&mut output)?;
462 }
463 output.push(
464 u32::try_from(part.extents.len())
465 .map_err(|_| PreparedInputError::WireValueOverflow("extent count"))?,
466 );
467 for extent in part.extents.values().copied() {
468 extent.encode_words(&mut output)?;
469 }
470 }
471 Ok(output)
472 }
473
474 pub fn decode_words(words: &[u32]) -> Result<Self, PreparedInputError> {
476 let mut cursor = WordCursor { words, offset: 0 };
477 let part_count = cursor.usize("part count")?;
478 if part_count == 0 {
479 return Err(PreparedInputError::EmptyInput);
480 }
481 let mut parts = Vec::with_capacity(part_count);
482 for _ in 0..part_count {
483 let modality = InputModality::from_wire_tag(cursor.next("modality")?)?;
484 let payload_kind = InputPayloadKind::from_wire_tag(cursor.next("payload kind")?)?;
485 let payload = InputTensorIdentity::decode_words(&mut cursor)?;
486 let metadata_count = cursor.usize("metadata count")?;
487 if metadata_count > 3 {
488 return Err(PreparedInputError::InvalidMetadataCount(metadata_count));
489 }
490 let metadata = (0..metadata_count)
491 .map(|_| {
492 let key = InputMetadataKey::from_wire_tag(cursor.next("metadata key")?)?;
493 Ok((key, InputTensorIdentity::decode_words(&mut cursor)?))
494 })
495 .collect::<Result<Vec<_>, PreparedInputError>>()?;
496 let extent_count = cursor.usize("extent count")?;
497 if extent_count > 2 {
498 return Err(PreparedInputError::InvalidExtentCount(extent_count));
499 }
500 let extents = (0..extent_count)
501 .map(|_| InputExtent::decode_words(&mut cursor))
502 .collect::<Result<Vec<_>, _>>()?;
503 parts.push(InputPartDescriptor::new_with_extents(
504 modality,
505 payload_kind,
506 payload,
507 metadata,
508 extents,
509 )?);
510 }
511 if cursor.offset != words.len() {
512 return Err(PreparedInputError::TrailingWireValues);
513 }
514 Self::new(parts)
515 }
516}
517
518#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
520#[non_exhaustive]
521pub enum PreparedInputError {
522 #[error("prepared model input must contain at least one part")]
524 EmptyInput,
525 #[error("prepared input tensor has invalid shape {shape:?}")]
527 InvalidTensorShape {
528 shape: Vec<usize>,
530 },
531 #[error("{modality:?} input is incompatible with {payload:?} payload")]
533 IncompatiblePayload {
534 modality: InputModality,
536 payload: InputPayloadKind,
538 },
539 #[error("{key:?} metadata is incompatible with {modality:?} input")]
541 IncompatibleMetadata {
542 modality: InputModality,
544 key: InputMetadataKey,
546 },
547 #[error("prepared input contains duplicate {key:?} metadata")]
549 DuplicateMetadata {
550 key: InputMetadataKey,
552 },
553 #[error("{extent:?} extent is incompatible with {modality:?} input")]
555 IncompatibleExtent {
556 modality: InputModality,
558 extent: InputExtent,
560 },
561 #[error("prepared input contains duplicate {extent:?} extent")]
563 DuplicateExtent {
564 extent: InputExtent,
566 },
567 #[error("prepared input part {part} ({modality:?}) is missing {key:?} metadata")]
569 MissingMetadata {
570 part: usize,
572 modality: InputModality,
574 key: InputMetadataKey,
576 },
577 #[error("prepared-input descriptor has invalid {field} value {value}")]
579 InvalidWireValue {
580 field: &'static str,
582 value: u32,
584 },
585 #[error("prepared-input descriptor ended while reading {0}")]
587 TruncatedWireDescriptor(&'static str),
588 #[error("prepared-input descriptor tensor rank {0} is outside 1..=8")]
590 InvalidWireRank(usize),
591 #[error("prepared-input descriptor metadata count {0} exceeds 3")]
593 InvalidMetadataCount(usize),
594 #[error("prepared-input descriptor extent count {0} exceeds 2")]
596 InvalidExtentCount(usize),
597 #[error("prepared-input descriptor has trailing values")]
599 TrailingWireValues,
600 #[error("prepared-input {0} exceeds descriptor range")]
602 WireValueOverflow(&'static str),
603 #[error("prepared-input wire payload has {actual} values; expected {expected}")]
605 WireValueCount {
606 expected: usize,
608 actual: usize,
610 },
611 #[error("prepared-input wire payload does not match its identity")]
613 WireIdentityMismatch,
614 #[error("encoded dtype {0:?} cannot identify a prepared runtime tensor")]
616 EncodedRuntimeDtype(String),
617 #[error("backend prepared-tensor identity failed: {0}")]
619 BackendTensorIdentity(String),
620}
621
622struct WordCursor<'a> {
623 words: &'a [u32],
624 offset: usize,
625}
626
627impl WordCursor<'_> {
628 fn next(&mut self, field: &'static str) -> Result<u32, PreparedInputError> {
629 let value = self
630 .words
631 .get(self.offset)
632 .copied()
633 .ok_or(PreparedInputError::TruncatedWireDescriptor(field))?;
634 self.offset += 1;
635 Ok(value)
636 }
637
638 fn usize(&mut self, field: &'static str) -> Result<usize, PreparedInputError> {
639 usize::try_from(self.next(field)?).map_err(|_| PreparedInputError::WireValueOverflow(field))
640 }
641}
642
643fn encode_dtype(dtype: &TensorDtype, output: &mut Vec<u32>) -> Result<(), PreparedInputError> {
644 let tag = match dtype {
645 TensorDtype::Bool => 0,
646 TensorDtype::U8 => 1,
647 TensorDtype::U16 => 2,
648 TensorDtype::U32 => 3,
649 TensorDtype::U64 => 4,
650 TensorDtype::I8 => 5,
651 TensorDtype::I16 => 6,
652 TensorDtype::I32 => 7,
653 TensorDtype::I64 => 8,
654 TensorDtype::F16 => 9,
655 TensorDtype::F32 => 10,
656 TensorDtype::F64 => 11,
657 TensorDtype::Bf16 => 12,
658 TensorDtype::Complex64 => 13,
659 TensorDtype::Encoded(name) => {
660 return Err(PreparedInputError::EncodedRuntimeDtype(name.clone()))
661 }
662 };
663 output.push(tag);
664 Ok(())
665}
666
667fn decode_dtype(cursor: &mut WordCursor<'_>) -> Result<TensorDtype, PreparedInputError> {
668 let tag = cursor.next("dtype")?;
669 match tag {
670 0 => Ok(TensorDtype::Bool),
671 1 => Ok(TensorDtype::U8),
672 2 => Ok(TensorDtype::U16),
673 3 => Ok(TensorDtype::U32),
674 4 => Ok(TensorDtype::U64),
675 5 => Ok(TensorDtype::I8),
676 6 => Ok(TensorDtype::I16),
677 7 => Ok(TensorDtype::I32),
678 8 => Ok(TensorDtype::I64),
679 9 => Ok(TensorDtype::F16),
680 10 => Ok(TensorDtype::F32),
681 11 => Ok(TensorDtype::F64),
682 12 => Ok(TensorDtype::Bf16),
683 13 => Ok(TensorDtype::Complex64),
684 value => Err(PreparedInputError::InvalidWireValue {
685 field: "dtype",
686 value,
687 }),
688 }
689}
690
691#[cfg(test)]
692mod tests {
693 use super::*;
694
695 fn tensor(dtype: TensorDtype, shape: &[usize]) -> InputTensorIdentity {
696 InputTensorIdentity::new(dtype, shape.to_vec()).unwrap()
697 }
698
699 #[test]
700 fn identity_round_trip_preserves_order_geometry_and_typed_metadata() {
701 let identity = PreparedInputIdentity::new(vec![
702 InputPartDescriptor::new(
703 InputModality::Text,
704 InputPayloadKind::TokenIds,
705 tensor(TensorDtype::U32, &[1, 2]),
706 [],
707 )
708 .unwrap(),
709 InputPartDescriptor::new_with_extents(
710 InputModality::Image,
711 InputPayloadKind::Tensor,
712 tensor(TensorDtype::F32, &[4, 12]),
713 [(
714 InputMetadataKey::PatchGrid,
715 tensor(TensorDtype::I32, &[1, 3]),
716 )],
717 [InputExtent::PatchGrid {
718 time: 1,
719 height: 2,
720 width: 2,
721 }],
722 )
723 .unwrap(),
724 ])
725 .unwrap();
726
727 let words = identity.encode_words().unwrap();
728 assert_eq!(
729 PreparedInputIdentity::decode_words(&words).unwrap(),
730 identity
731 );
732 assert_eq!(
733 identity.parts()[1].extents().collect::<Vec<_>>(),
734 [InputExtent::PatchGrid {
735 time: 1,
736 height: 2,
737 width: 2,
738 }]
739 );
740 }
741
742 #[test]
743 fn rejects_duplicate_missing_and_modality_incompatible_metadata() {
744 let grid = tensor(TensorDtype::I32, &[1, 3]);
745 let duplicate = InputPartDescriptor::new(
746 InputModality::Image,
747 InputPayloadKind::Tensor,
748 tensor(TensorDtype::F32, &[2, 4]),
749 [
750 (InputMetadataKey::PatchGrid, grid.clone()),
751 (InputMetadataKey::PatchGrid, grid.clone()),
752 ],
753 );
754 assert!(matches!(
755 duplicate,
756 Err(PreparedInputError::DuplicateMetadata { .. })
757 ));
758
759 let image = InputPartDescriptor::new(
760 InputModality::Image,
761 InputPayloadKind::Tensor,
762 tensor(TensorDtype::F32, &[2, 4]),
763 [],
764 )
765 .unwrap();
766 assert!(matches!(
767 image.require_metadata(0, InputMetadataKey::PatchGrid),
768 Err(PreparedInputError::MissingMetadata { .. })
769 ));
770
771 assert!(matches!(
772 InputPartDescriptor::new(
773 InputModality::Audio,
774 InputPayloadKind::Tensor,
775 tensor(TensorDtype::F32, &[2, 4]),
776 [(InputMetadataKey::PatchGrid, grid)]
777 ),
778 Err(PreparedInputError::IncompatibleMetadata { .. })
779 ));
780 }
781
782 #[test]
783 fn malformed_wire_descriptors_fail_closed() {
784 assert!(matches!(
785 PreparedInputIdentity::decode_words(&[]),
786 Err(PreparedInputError::TruncatedWireDescriptor(_))
787 ));
788 assert!(matches!(
789 PreparedInputIdentity::decode_words(&[1, 99]),
790 Err(PreparedInputError::InvalidWireValue {
791 field: "modality",
792 ..
793 })
794 ));
795 }
796}