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    /// Logical host storage for this payload-free description, including shapes,
406    /// map entries and encoded dtype names. Excludes allocator capacity/overhead;
407    /// no native tensor is inspected or allocated. Overflow returns `None`.
408    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    /// Number of ordered parts.
436    pub fn len(&self) -> usize {
437        self.parts.len()
438    }
439
440    /// This identity is always non-empty after construction.
441    pub fn is_empty(&self) -> bool {
442        self.parts.is_empty()
443    }
444
445    /// Encodes the identity for backend-independent rank agreement.
446    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    /// Decodes and validates a rank-agreement descriptor.
475    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/// Invalid portable prepared-input identity.
519#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
520#[non_exhaustive]
521pub enum PreparedInputError {
522    /// At least one ordered part is required.
523    #[error("prepared model input must contain at least one part")]
524    EmptyInput,
525    /// Tensor identities must be non-scalar and have non-zero dimensions.
526    #[error("prepared input tensor has invalid shape {shape:?}")]
527    InvalidTensorShape {
528        /// Invalid logical dimensions.
529        shape: Vec<usize>,
530    },
531    /// A modality cannot consume the selected primary payload kind.
532    #[error("{modality:?} input is incompatible with {payload:?} payload")]
533    IncompatiblePayload {
534        /// Declared modality.
535        modality: InputModality,
536        /// Declared payload kind.
537        payload: InputPayloadKind,
538    },
539    /// Metadata belongs to a different modality.
540    #[error("{key:?} metadata is incompatible with {modality:?} input")]
541    IncompatibleMetadata {
542        /// Declared modality.
543        modality: InputModality,
544        /// Metadata key.
545        key: InputMetadataKey,
546    },
547    /// A typed metadata key occurred more than once.
548    #[error("prepared input contains duplicate {key:?} metadata")]
549    DuplicateMetadata {
550        /// Duplicated key.
551        key: InputMetadataKey,
552    },
553    /// A host execution extent belongs to a different modality.
554    #[error("{extent:?} extent is incompatible with {modality:?} input")]
555    IncompatibleExtent {
556        /// Declared modality.
557        modality: InputModality,
558        /// Incompatible extent.
559        extent: InputExtent,
560    },
561    /// An execution extent kind occurred more than once.
562    #[error("prepared input contains duplicate {extent:?} extent")]
563    DuplicateExtent {
564        /// Duplicated extent.
565        extent: InputExtent,
566    },
567    /// Family policy required absent metadata.
568    #[error("prepared input part {part} ({modality:?}) is missing {key:?} metadata")]
569    MissingMetadata {
570        /// Ordered part index.
571        part: usize,
572        /// Part modality.
573        modality: InputModality,
574        /// Required key.
575        key: InputMetadataKey,
576    },
577    /// A descriptor tag was not recognized.
578    #[error("prepared-input descriptor has invalid {field} value {value}")]
579    InvalidWireValue {
580        /// Descriptor field.
581        field: &'static str,
582        /// Invalid value.
583        value: u32,
584    },
585    /// Descriptor ended before the announced data was present.
586    #[error("prepared-input descriptor ended while reading {0}")]
587    TruncatedWireDescriptor(&'static str),
588    /// Descriptor carried an unsupported tensor rank.
589    #[error("prepared-input descriptor tensor rank {0} is outside 1..=8")]
590    InvalidWireRank(usize),
591    /// A part advertised more metadata keys than the closed vocabulary contains.
592    #[error("prepared-input descriptor metadata count {0} exceeds 3")]
593    InvalidMetadataCount(usize),
594    /// A part advertised more host extents than the closed vocabulary contains.
595    #[error("prepared-input descriptor extent count {0} exceeds 2")]
596    InvalidExtentCount(usize),
597    /// Descriptor has data after the complete identity.
598    #[error("prepared-input descriptor has trailing values")]
599    TrailingWireValues,
600    /// A host value cannot be represented by the u32 wire format.
601    #[error("prepared-input {0} exceeds descriptor range")]
602    WireValueOverflow(&'static str),
603    /// The number of received tensors does not match the identity.
604    #[error("prepared-input wire payload has {actual} values; expected {expected}")]
605    WireValueCount {
606        /// Expected payload and metadata tensor count.
607        expected: usize,
608        /// Received tensor count.
609        actual: usize,
610    },
611    /// Received tensor shape or dtype disagrees with the advertised identity.
612    #[error("prepared-input wire payload does not match its identity")]
613    WireIdentityMismatch,
614    /// Encoded/quantized checkpoint storage is not a runtime tensor dtype.
615    #[error("encoded dtype {0:?} cannot identify a prepared runtime tensor")]
616    EncodedRuntimeDtype(String),
617    /// A backend could not describe a prepared tensor without materializing it.
618    #[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}