Skip to main content

eredu_core/
input.rs

1//! Portable identity for ordered, prepared model input.
2
3use std::collections::BTreeMap;
4
5use serde::{Deserialize, Serialize};
6
7use crate::checkpoint::TensorDtype;
8
9/// Modality of one ordered model-input part.
10#[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 token IDs or precomputed text embeddings.
15    Text,
16    /// Still-image patches or precomputed image embeddings.
17    Image,
18    /// Video patches or precomputed video embeddings.
19    Video,
20    /// Audio features or precomputed audio embeddings.
21    Audio,
22}
23
24impl InputModality {
25    /// Stable lowercase diagnostic name.
26    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/// Semantic role of a prepared part's primary tensor.
59#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
60#[serde(rename_all = "snake_case")]
61#[non_exhaustive]
62pub enum InputPayloadKind {
63    /// Tokenizer vocabulary IDs.
64    TokenIds,
65    /// Model-native media features or patches that still require an encoder.
66    Tensor,
67    /// Already projected decoder-width embeddings.
68    Embeddings,
69}
70
71impl InputPayloadKind {
72    /// Returns whether this payload role is meaningful for `modality`.
73    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/// Architecture-neutral metadata attached to a prepared input part.
103#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
104#[serde(rename_all = "snake_case")]
105#[non_exhaustive]
106pub enum InputMetadataKey {
107    /// One or more `(time, height, width)` patch-grid rows.
108    PatchGrid,
109    /// Explicit spatial or temporal patch coordinates, including padding.
110    PatchPositions,
111    /// Valid-frame or valid-feature mask for audio.
112    AudioMask,
113}
114
115/// Host-known extent needed to plan device execution without reading a tensor
116/// back from the accelerator.
117#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
118#[serde(rename_all = "snake_case")]
119#[non_exhaustive]
120pub enum InputExtent {
121    /// Exact `(time, height, width)` of one image or video patch grid.
122    PatchGrid {
123        /// Temporal extent (one for an independently packed frame).
124        time: usize,
125        /// Unpadded patch-row count.
126        height: usize,
127        /// Unpadded patch-column count.
128        width: usize,
129    },
130    /// Number of valid (unpadded) input audio frames.
131    AudioValidFrames(usize),
132}
133
134impl InputExtent {
135    /// Returns whether this extent is meaningful for `modality`.
136    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    /// Returns whether this metadata is meaningful for `modality`.
193    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/// Payload-free shape and element-type identity for a backend tensor.
224#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
225pub struct InputTensorIdentity {
226    dtype: TensorDtype,
227    shape: Vec<usize>,
228}
229
230impl InputTensorIdentity {
231    /// Validates a non-scalar tensor identity with non-zero dimensions.
232    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    /// Logical element type.
240    pub const fn dtype(&self) -> &TensorDtype {
241        &self.dtype
242    }
243
244    /// Row-major logical shape.
245    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/// Payload-free identity for one ordered prepared input part.
278#[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    /// Validates modality/payload compatibility and unique typed metadata.
289    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    /// Validates tensor metadata plus host-known execution extents.
299    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    /// Part modality.
340    pub const fn modality(&self) -> InputModality {
341        self.modality
342    }
343
344    /// Primary tensor role.
345    pub const fn payload_kind(&self) -> InputPayloadKind {
346        self.payload_kind
347    }
348
349    /// Primary tensor identity.
350    pub const fn payload(&self) -> &InputTensorIdentity {
351        &self.payload
352    }
353
354    /// Typed metadata identities in stable key order.
355    pub const fn metadata(&self) -> &BTreeMap<InputMetadataKey, InputTensorIdentity> {
356        &self.metadata
357    }
358
359    /// Host-known execution extents in stable wire order.
360    pub fn extents(&self) -> impl ExactSizeIterator<Item = InputExtent> + '_ {
361        self.extents.values().copied()
362    }
363
364    /// Looks up one typed metadata identity.
365    pub fn metadata_value(&self, key: InputMetadataKey) -> Option<&InputTensorIdentity> {
366        self.metadata.get(&key)
367    }
368
369    /// Requires metadata selected by family policy.
370    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/// Exact payload-free identity of an ordered prepared model input.
386#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
387pub struct PreparedInputIdentity {
388    parts: Vec<InputPartDescriptor>,
389}
390
391impl PreparedInputIdentity {
392    /// Validates a non-empty ordered set of input-part descriptors.
393    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    /// Ordered input-part descriptors.
401    pub fn parts(&self) -> &[InputPartDescriptor] {
402        &self.parts
403    }
404
405    /// Number of ordered parts.
406    pub fn len(&self) -> usize {
407        self.parts.len()
408    }
409
410    /// This identity is always non-empty after construction.
411    pub fn is_empty(&self) -> bool {
412        self.parts.is_empty()
413    }
414
415    /// Encodes the identity for backend-independent rank agreement.
416    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    /// Decodes and validates a rank-agreement descriptor.
445    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/// Invalid portable prepared-input identity.
489#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
490#[non_exhaustive]
491pub enum PreparedInputError {
492    /// At least one ordered part is required.
493    #[error("prepared model input must contain at least one part")]
494    EmptyInput,
495    /// Tensor identities must be non-scalar and have non-zero dimensions.
496    #[error("prepared input tensor has invalid shape {shape:?}")]
497    InvalidTensorShape {
498        /// Invalid logical dimensions.
499        shape: Vec<usize>,
500    },
501    /// A modality cannot consume the selected primary payload kind.
502    #[error("{modality:?} input is incompatible with {payload:?} payload")]
503    IncompatiblePayload {
504        /// Declared modality.
505        modality: InputModality,
506        /// Declared payload kind.
507        payload: InputPayloadKind,
508    },
509    /// Metadata belongs to a different modality.
510    #[error("{key:?} metadata is incompatible with {modality:?} input")]
511    IncompatibleMetadata {
512        /// Declared modality.
513        modality: InputModality,
514        /// Metadata key.
515        key: InputMetadataKey,
516    },
517    /// A typed metadata key occurred more than once.
518    #[error("prepared input contains duplicate {key:?} metadata")]
519    DuplicateMetadata {
520        /// Duplicated key.
521        key: InputMetadataKey,
522    },
523    /// A host execution extent belongs to a different modality.
524    #[error("{extent:?} extent is incompatible with {modality:?} input")]
525    IncompatibleExtent {
526        /// Declared modality.
527        modality: InputModality,
528        /// Incompatible extent.
529        extent: InputExtent,
530    },
531    /// An execution extent kind occurred more than once.
532    #[error("prepared input contains duplicate {extent:?} extent")]
533    DuplicateExtent {
534        /// Duplicated extent.
535        extent: InputExtent,
536    },
537    /// Family policy required absent metadata.
538    #[error("prepared input part {part} ({modality:?}) is missing {key:?} metadata")]
539    MissingMetadata {
540        /// Ordered part index.
541        part: usize,
542        /// Part modality.
543        modality: InputModality,
544        /// Required key.
545        key: InputMetadataKey,
546    },
547    /// A descriptor tag was not recognized.
548    #[error("prepared-input descriptor has invalid {field} value {value}")]
549    InvalidWireValue {
550        /// Descriptor field.
551        field: &'static str,
552        /// Invalid value.
553        value: u32,
554    },
555    /// Descriptor ended before the announced data was present.
556    #[error("prepared-input descriptor ended while reading {0}")]
557    TruncatedWireDescriptor(&'static str),
558    /// Descriptor carried an unsupported tensor rank.
559    #[error("prepared-input descriptor tensor rank {0} is outside 1..=8")]
560    InvalidWireRank(usize),
561    /// A part advertised more metadata keys than the closed vocabulary contains.
562    #[error("prepared-input descriptor metadata count {0} exceeds 3")]
563    InvalidMetadataCount(usize),
564    /// A part advertised more host extents than the closed vocabulary contains.
565    #[error("prepared-input descriptor extent count {0} exceeds 2")]
566    InvalidExtentCount(usize),
567    /// Descriptor has data after the complete identity.
568    #[error("prepared-input descriptor has trailing values")]
569    TrailingWireValues,
570    /// A host value cannot be represented by the u32 wire format.
571    #[error("prepared-input {0} exceeds descriptor range")]
572    WireValueOverflow(&'static str),
573    /// The number of received tensors does not match the identity.
574    #[error("prepared-input wire payload has {actual} values; expected {expected}")]
575    WireValueCount {
576        /// Expected payload and metadata tensor count.
577        expected: usize,
578        /// Received tensor count.
579        actual: usize,
580    },
581    /// Received tensor shape or dtype disagrees with the advertised identity.
582    #[error("prepared-input wire payload does not match its identity")]
583    WireIdentityMismatch,
584    /// Encoded/quantized checkpoint storage is not a runtime tensor dtype.
585    #[error("encoded dtype {0:?} cannot identify a prepared runtime tensor")]
586    EncodedRuntimeDtype(String),
587    /// A backend could not describe a prepared tensor without materializing it.
588    #[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}