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 len(&self) -> usize {
407 self.parts.len()
408 }
409
410 pub fn is_empty(&self) -> bool {
412 self.parts.is_empty()
413 }
414
415 pub fn encode_words(&self) -> Result<Vec<u32>, PreparedInputError> {
417 let mut output = Vec::new();
418 output.push(
419 u32::try_from(self.parts.len())
420 .map_err(|_| PreparedInputError::WireValueOverflow("part count"))?,
421 );
422 for part in &self.parts {
423 output.extend_from_slice(&[part.modality.wire_tag(), part.payload_kind.wire_tag()]);
424 part.payload.encode_words(&mut output)?;
425 output.push(
426 u32::try_from(part.metadata.len())
427 .map_err(|_| PreparedInputError::WireValueOverflow("metadata count"))?,
428 );
429 for (key, identity) in &part.metadata {
430 output.push(key.wire_tag());
431 identity.encode_words(&mut output)?;
432 }
433 output.push(
434 u32::try_from(part.extents.len())
435 .map_err(|_| PreparedInputError::WireValueOverflow("extent count"))?,
436 );
437 for extent in part.extents.values().copied() {
438 extent.encode_words(&mut output)?;
439 }
440 }
441 Ok(output)
442 }
443
444 pub fn decode_words(words: &[u32]) -> Result<Self, PreparedInputError> {
446 let mut cursor = WordCursor { words, offset: 0 };
447 let part_count = cursor.usize("part count")?;
448 if part_count == 0 {
449 return Err(PreparedInputError::EmptyInput);
450 }
451 let mut parts = Vec::with_capacity(part_count);
452 for _ in 0..part_count {
453 let modality = InputModality::from_wire_tag(cursor.next("modality")?)?;
454 let payload_kind = InputPayloadKind::from_wire_tag(cursor.next("payload kind")?)?;
455 let payload = InputTensorIdentity::decode_words(&mut cursor)?;
456 let metadata_count = cursor.usize("metadata count")?;
457 if metadata_count > 3 {
458 return Err(PreparedInputError::InvalidMetadataCount(metadata_count));
459 }
460 let metadata = (0..metadata_count)
461 .map(|_| {
462 let key = InputMetadataKey::from_wire_tag(cursor.next("metadata key")?)?;
463 Ok((key, InputTensorIdentity::decode_words(&mut cursor)?))
464 })
465 .collect::<Result<Vec<_>, PreparedInputError>>()?;
466 let extent_count = cursor.usize("extent count")?;
467 if extent_count > 2 {
468 return Err(PreparedInputError::InvalidExtentCount(extent_count));
469 }
470 let extents = (0..extent_count)
471 .map(|_| InputExtent::decode_words(&mut cursor))
472 .collect::<Result<Vec<_>, _>>()?;
473 parts.push(InputPartDescriptor::new_with_extents(
474 modality,
475 payload_kind,
476 payload,
477 metadata,
478 extents,
479 )?);
480 }
481 if cursor.offset != words.len() {
482 return Err(PreparedInputError::TrailingWireValues);
483 }
484 Self::new(parts)
485 }
486}
487
488#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
490#[non_exhaustive]
491pub enum PreparedInputError {
492 #[error("prepared model input must contain at least one part")]
494 EmptyInput,
495 #[error("prepared input tensor has invalid shape {shape:?}")]
497 InvalidTensorShape {
498 shape: Vec<usize>,
500 },
501 #[error("{modality:?} input is incompatible with {payload:?} payload")]
503 IncompatiblePayload {
504 modality: InputModality,
506 payload: InputPayloadKind,
508 },
509 #[error("{key:?} metadata is incompatible with {modality:?} input")]
511 IncompatibleMetadata {
512 modality: InputModality,
514 key: InputMetadataKey,
516 },
517 #[error("prepared input contains duplicate {key:?} metadata")]
519 DuplicateMetadata {
520 key: InputMetadataKey,
522 },
523 #[error("{extent:?} extent is incompatible with {modality:?} input")]
525 IncompatibleExtent {
526 modality: InputModality,
528 extent: InputExtent,
530 },
531 #[error("prepared input contains duplicate {extent:?} extent")]
533 DuplicateExtent {
534 extent: InputExtent,
536 },
537 #[error("prepared input part {part} ({modality:?}) is missing {key:?} metadata")]
539 MissingMetadata {
540 part: usize,
542 modality: InputModality,
544 key: InputMetadataKey,
546 },
547 #[error("prepared-input descriptor has invalid {field} value {value}")]
549 InvalidWireValue {
550 field: &'static str,
552 value: u32,
554 },
555 #[error("prepared-input descriptor ended while reading {0}")]
557 TruncatedWireDescriptor(&'static str),
558 #[error("prepared-input descriptor tensor rank {0} is outside 1..=8")]
560 InvalidWireRank(usize),
561 #[error("prepared-input descriptor metadata count {0} exceeds 3")]
563 InvalidMetadataCount(usize),
564 #[error("prepared-input descriptor extent count {0} exceeds 2")]
566 InvalidExtentCount(usize),
567 #[error("prepared-input descriptor has trailing values")]
569 TrailingWireValues,
570 #[error("prepared-input {0} exceeds descriptor range")]
572 WireValueOverflow(&'static str),
573 #[error("prepared-input wire payload has {actual} values; expected {expected}")]
575 WireValueCount {
576 expected: usize,
578 actual: usize,
580 },
581 #[error("prepared-input wire payload does not match its identity")]
583 WireIdentityMismatch,
584 #[error("encoded dtype {0:?} cannot identify a prepared runtime tensor")]
586 EncodedRuntimeDtype(String),
587 #[error("backend prepared-tensor identity failed: {0}")]
589 BackendTensorIdentity(String),
590}
591
592struct WordCursor<'a> {
593 words: &'a [u32],
594 offset: usize,
595}
596
597impl WordCursor<'_> {
598 fn next(&mut self, field: &'static str) -> Result<u32, PreparedInputError> {
599 let value = self
600 .words
601 .get(self.offset)
602 .copied()
603 .ok_or(PreparedInputError::TruncatedWireDescriptor(field))?;
604 self.offset += 1;
605 Ok(value)
606 }
607
608 fn usize(&mut self, field: &'static str) -> Result<usize, PreparedInputError> {
609 usize::try_from(self.next(field)?).map_err(|_| PreparedInputError::WireValueOverflow(field))
610 }
611}
612
613fn encode_dtype(dtype: &TensorDtype, output: &mut Vec<u32>) -> Result<(), PreparedInputError> {
614 let tag = match dtype {
615 TensorDtype::Bool => 0,
616 TensorDtype::U8 => 1,
617 TensorDtype::U16 => 2,
618 TensorDtype::U32 => 3,
619 TensorDtype::U64 => 4,
620 TensorDtype::I8 => 5,
621 TensorDtype::I16 => 6,
622 TensorDtype::I32 => 7,
623 TensorDtype::I64 => 8,
624 TensorDtype::F16 => 9,
625 TensorDtype::F32 => 10,
626 TensorDtype::F64 => 11,
627 TensorDtype::Bf16 => 12,
628 TensorDtype::Complex64 => 13,
629 TensorDtype::Encoded(name) => {
630 return Err(PreparedInputError::EncodedRuntimeDtype(name.clone()))
631 }
632 };
633 output.push(tag);
634 Ok(())
635}
636
637fn decode_dtype(cursor: &mut WordCursor<'_>) -> Result<TensorDtype, PreparedInputError> {
638 let tag = cursor.next("dtype")?;
639 match tag {
640 0 => Ok(TensorDtype::Bool),
641 1 => Ok(TensorDtype::U8),
642 2 => Ok(TensorDtype::U16),
643 3 => Ok(TensorDtype::U32),
644 4 => Ok(TensorDtype::U64),
645 5 => Ok(TensorDtype::I8),
646 6 => Ok(TensorDtype::I16),
647 7 => Ok(TensorDtype::I32),
648 8 => Ok(TensorDtype::I64),
649 9 => Ok(TensorDtype::F16),
650 10 => Ok(TensorDtype::F32),
651 11 => Ok(TensorDtype::F64),
652 12 => Ok(TensorDtype::Bf16),
653 13 => Ok(TensorDtype::Complex64),
654 value => Err(PreparedInputError::InvalidWireValue {
655 field: "dtype",
656 value,
657 }),
658 }
659}
660
661#[cfg(test)]
662mod tests {
663 use super::*;
664
665 fn tensor(dtype: TensorDtype, shape: &[usize]) -> InputTensorIdentity {
666 InputTensorIdentity::new(dtype, shape.to_vec()).unwrap()
667 }
668
669 #[test]
670 fn identity_round_trip_preserves_order_geometry_and_typed_metadata() {
671 let identity = PreparedInputIdentity::new(vec![
672 InputPartDescriptor::new(
673 InputModality::Text,
674 InputPayloadKind::TokenIds,
675 tensor(TensorDtype::U32, &[1, 2]),
676 [],
677 )
678 .unwrap(),
679 InputPartDescriptor::new_with_extents(
680 InputModality::Image,
681 InputPayloadKind::Tensor,
682 tensor(TensorDtype::F32, &[4, 12]),
683 [(
684 InputMetadataKey::PatchGrid,
685 tensor(TensorDtype::I32, &[1, 3]),
686 )],
687 [InputExtent::PatchGrid {
688 time: 1,
689 height: 2,
690 width: 2,
691 }],
692 )
693 .unwrap(),
694 ])
695 .unwrap();
696
697 let words = identity.encode_words().unwrap();
698 assert_eq!(
699 PreparedInputIdentity::decode_words(&words).unwrap(),
700 identity
701 );
702 assert_eq!(
703 identity.parts()[1].extents().collect::<Vec<_>>(),
704 [InputExtent::PatchGrid {
705 time: 1,
706 height: 2,
707 width: 2,
708 }]
709 );
710 }
711
712 #[test]
713 fn rejects_duplicate_missing_and_modality_incompatible_metadata() {
714 let grid = tensor(TensorDtype::I32, &[1, 3]);
715 let duplicate = InputPartDescriptor::new(
716 InputModality::Image,
717 InputPayloadKind::Tensor,
718 tensor(TensorDtype::F32, &[2, 4]),
719 [
720 (InputMetadataKey::PatchGrid, grid.clone()),
721 (InputMetadataKey::PatchGrid, grid.clone()),
722 ],
723 );
724 assert!(matches!(
725 duplicate,
726 Err(PreparedInputError::DuplicateMetadata { .. })
727 ));
728
729 let image = InputPartDescriptor::new(
730 InputModality::Image,
731 InputPayloadKind::Tensor,
732 tensor(TensorDtype::F32, &[2, 4]),
733 [],
734 )
735 .unwrap();
736 assert!(matches!(
737 image.require_metadata(0, InputMetadataKey::PatchGrid),
738 Err(PreparedInputError::MissingMetadata { .. })
739 ));
740
741 assert!(matches!(
742 InputPartDescriptor::new(
743 InputModality::Audio,
744 InputPayloadKind::Tensor,
745 tensor(TensorDtype::F32, &[2, 4]),
746 [(InputMetadataKey::PatchGrid, grid)]
747 ),
748 Err(PreparedInputError::IncompatibleMetadata { .. })
749 ));
750 }
751
752 #[test]
753 fn malformed_wire_descriptors_fail_closed() {
754 assert!(matches!(
755 PreparedInputIdentity::decode_words(&[]),
756 Err(PreparedInputError::TruncatedWireDescriptor(_))
757 ));
758 assert!(matches!(
759 PreparedInputIdentity::decode_words(&[1, 99]),
760 Err(PreparedInputError::InvalidWireValue {
761 field: "modality",
762 ..
763 })
764 ));
765 }
766}