tract_api/
lib.rs

1use anyhow::{Result, ensure};
2use boow::Bow;
3use std::fmt::{Debug, Display};
4use std::path::Path;
5
6#[macro_use]
7pub mod macros;
8pub mod transform;
9
10pub use transform::{FloatPrecision, Pulse, SetSymbols, TransformConfig, TransformSpec};
11
12/// an implementation of tract's NNEF framework object
13///
14/// Entry point for NNEF model manipulation: loading from file, dumping to file.
15pub trait NnefInterface: Debug + Sized {
16    type Model: ModelInterface;
17    /// Load a NNEF model from the path into a tract-core model.
18    ///
19    /// * `path` can point to a directory, a `tar` file or a `tar.gz` file.
20    fn load(&self, path: impl AsRef<Path>) -> Result<Self::Model>;
21
22    /// Load a NNEF model from a buffer into a tract-core model.
23    ///
24    /// data is the content of a NNEF model, as a `tar` file or a `tar.gz` file.
25    fn load_buffer(&self, data: &[u8]) -> Result<Self::Model>;
26
27    /// Force the framework to emit strict NNEF instead of using the tract_core extension.
28    /// The tract_core extension is enabled by default; call this to opt out.
29    fn disable_tract_core(&mut self) -> Result<()>;
30
31    /// Allow the framework to use tract_extra extensions.
32    fn enable_tract_extra(&mut self) -> Result<()>;
33
34    /// Allow the framework to use tract_transformers extensions to support common transformer operators.
35    fn enable_tract_transformers(&mut self) -> Result<()>;
36
37    /// Allow the framework to use tract_onnx extensions to support operators in ONNX that are
38    /// absent from NNEF.
39    fn enable_onnx(&mut self) -> Result<()>;
40
41    /// Allow the framework to use tract_pulse extensions to support stateful streaming operation.
42    fn enable_pulse(&mut self) -> Result<()>;
43
44    /// Allow the framework to use a tract-proprietary extension that can support special characters
45    /// in node names. If disable, tract will replace everything by underscore '_' to keep
46    /// compatibility with NNEF. If enabled, the extended syntax will be used, allowing to maintain
47    /// the node names in serialized form.
48    fn enable_extended_identifier_syntax(&mut self) -> Result<()>;
49
50    /// Convenience function, similar to disable_tract_core but allowing method chaining.
51    fn without_tract_core(mut self) -> Result<Self> {
52        self.disable_tract_core()?;
53        Ok(self)
54    }
55
56    /// Convenience function, similar with enable_tract_extra but allowing method chaining.
57    fn with_tract_extra(mut self) -> Result<Self> {
58        self.enable_tract_extra()?;
59        Ok(self)
60    }
61
62    /// Convenience function, similar with enable_tract_transformers but allowing method chaining.
63    fn with_tract_transformers(mut self) -> Result<Self> {
64        self.enable_tract_transformers()?;
65        Ok(self)
66    }
67
68    /// Convenience function, similar with enable_onnx but allowing method chaining.
69    fn with_onnx(mut self) -> Result<Self> {
70        self.enable_onnx()?;
71        Ok(self)
72    }
73
74    /// Convenience function, similar with enable_pulse but allowing method chaining.
75    fn with_pulse(mut self) -> Result<Self> {
76        self.enable_pulse()?;
77        Ok(self)
78    }
79
80    /// Convenience function, similar with enable_extended_identifier_syntax but allowing method chaining.
81    fn with_extended_identifier_syntax(mut self) -> Result<Self> {
82        self.enable_extended_identifier_syntax()?;
83        Ok(self)
84    }
85
86    /// Dump a TypedModel as a NNEF directory.
87    ///
88    /// `path` is the directory name to dump to
89    fn write_model_to_dir(&self, path: impl AsRef<Path>, model: &Self::Model) -> Result<()>;
90
91    /// Dump a TypedModel as a NNEF tar file.
92    ///
93    /// This function creates a plain, non-compressed, archive.
94    ///
95    /// `path` is the archive name
96    fn write_model_to_tar(&self, path: impl AsRef<Path>, model: &Self::Model) -> Result<()>;
97    fn write_model_to_tar_gz(&self, path: impl AsRef<Path>, model: &Self::Model) -> Result<()>;
98}
99
100/// Options for the ONNX loader, carried as JSON.
101///
102/// Follows the same principle as [`TransformSpec`]: a typed struct that
103/// serializes to a JSON object, so a caller can pass either the struct or a
104/// JSON string, and the C and Python entry points can pass the string.
105///
106/// An unknown field is an error rather than being ignored, so a misspelling
107/// cannot look like an option that took effect.
108#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
109#[serde(default, deny_unknown_fields)]
110pub struct OnnxOptions {
111    /// Discard the shapes the model file declares and let tract infer them.
112    pub ignore_value_info: bool,
113    /// Assertions over the model symbols, applied once the model is parsed, for
114    /// instance `["h>=0", "w>=0"]`. Each one goes to tract's assertion parser
115    /// unchanged.
116    pub assertions: Vec<String>,
117}
118
119impl OnnxOptions {
120    pub fn new() -> Self {
121        Self::default()
122    }
123
124    /// Discard the shapes the model file declares and let tract infer them.
125    pub fn ignore_value_info(mut self, ignore: bool) -> Self {
126        self.ignore_value_info = ignore;
127        self
128    }
129
130    /// Add one assertion over the model symbols.
131    pub fn assertion(mut self, assertion: impl Into<String>) -> Self {
132        self.assertions.push(assertion.into());
133        self
134    }
135
136    /// Produce the JSON the loader expects.
137    pub fn to_json(&self) -> String {
138        serde_json::to_string(self).expect("OnnxOptions serialization cannot fail")
139    }
140}
141
142/// A serialized [`OnnxOptions`], accepted as either the struct or a JSON string.
143#[derive(Debug, Clone)]
144pub struct OnnxOptionsSpec(String);
145
146impl OnnxOptionsSpec {
147    /// The JSON the loader expects.
148    pub fn to_json(&self) -> String {
149        self.0.clone()
150    }
151
152    /// Parse back into the typed form.
153    pub fn parse(&self) -> Result<OnnxOptions> {
154        Ok(serde_json::from_str(&self.0)?)
155    }
156}
157
158impl From<OnnxOptions> for OnnxOptionsSpec {
159    fn from(o: OnnxOptions) -> Self {
160        OnnxOptionsSpec(o.to_json())
161    }
162}
163
164impl From<&str> for OnnxOptionsSpec {
165    fn from(s: &str) -> Self {
166        OnnxOptionsSpec(s.to_string())
167    }
168}
169
170impl From<String> for OnnxOptionsSpec {
171    fn from(s: String) -> Self {
172        OnnxOptionsSpec(s)
173    }
174}
175
176pub trait OnnxInterface: Debug {
177    type InferenceModel: InferenceModelInterface;
178
179    /// Load a ONNX model from a file into an InferenceModel.
180    fn load(&self, path: impl AsRef<Path>) -> Result<Self::InferenceModel> {
181        self.load_with_options(path, OnnxOptions::new())
182    }
183
184    /// Load a ONNX model from a buffer into an InferenceModel.
185    fn load_buffer(&self, data: &[u8]) -> Result<Self::InferenceModel> {
186        self.load_buffer_with_options(data, OnnxOptions::new())
187    }
188
189    /// Load a ONNX model from a file, configuring the loader with `options`.
190    ///
191    /// `options` is either an [`OnnxOptions`] or the JSON for one.
192    fn load_with_options(
193        &self,
194        path: impl AsRef<Path>,
195        options: impl Into<OnnxOptionsSpec>,
196    ) -> Result<Self::InferenceModel>;
197
198    /// Load a ONNX model from a buffer, configuring the loader with `options`.
199    ///
200    /// See [`OnnxInterface::load_with_options`].
201    fn load_buffer_with_options(
202        &self,
203        data: &[u8],
204        options: impl Into<OnnxOptionsSpec>,
205    ) -> Result<Self::InferenceModel>;
206}
207
208pub trait InferenceModelInterface: Debug + Sized {
209    type Model: ModelInterface;
210    type InferenceFact: InferenceFactInterface;
211    fn input_count(&self) -> Result<usize>;
212    fn output_count(&self) -> Result<usize>;
213    fn input_name(&self, id: usize) -> Result<String>;
214    fn output_name(&self, id: usize) -> Result<String>;
215
216    fn input_fact(&self, id: usize) -> Result<Self::InferenceFact>;
217
218    fn set_input_fact(
219        &mut self,
220        id: usize,
221        fact: impl AsFact<Self, Self::InferenceFact>,
222    ) -> Result<()>;
223
224    fn output_fact(&self, id: usize) -> Result<Self::InferenceFact>;
225
226    fn set_output_fact(
227        &mut self,
228        id: usize,
229        fact: impl AsFact<Self, Self::InferenceFact>,
230    ) -> Result<()>;
231
232    fn analyse(&mut self) -> Result<()>;
233
234    fn into_model(self) -> Result<Self::Model>;
235}
236
237pub trait ModelInterface: Debug + Sized {
238    type Fact: FactInterface;
239    type Runnable: RunnableInterface;
240    type Tensor: TensorInterface;
241    fn input_count(&self) -> Result<usize>;
242
243    fn output_count(&self) -> Result<usize>;
244
245    fn input_name(&self, id: usize) -> Result<String>;
246
247    fn output_name(&self, id: usize) -> Result<String>;
248
249    fn input_fact(&self, id: usize) -> Result<Self::Fact>;
250
251    fn output_fact(&self, id: usize) -> Result<Self::Fact>;
252
253    fn into_runnable(self) -> Result<Self::Runnable>;
254
255    fn transform(&mut self, spec: impl Into<TransformSpec>) -> Result<()>;
256
257    fn property_keys(&self) -> Result<Vec<String>>;
258
259    fn property(&self, name: impl AsRef<str>) -> Result<Self::Tensor>;
260
261    fn parse_fact(&self, spec: &str) -> Result<Self::Fact>;
262
263    fn input_facts(&self) -> Result<impl Iterator<Item = Self::Fact>> {
264        Ok((0..self.input_count()?)
265            .map(|ix| self.input_fact(ix))
266            .collect::<Result<Vec<_>>>()?
267            .into_iter())
268    }
269
270    fn output_facts(&self) -> Result<impl Iterator<Item = Self::Fact>> {
271        Ok((0..self.output_count()?)
272            .map(|ix| self.output_fact(ix))
273            .collect::<Result<Vec<_>>>()?
274            .into_iter())
275    }
276}
277
278pub trait RuntimeInterface: Debug {
279    type Runnable: RunnableInterface;
280    type Model: ModelInterface;
281    fn name(&self) -> Result<String>;
282    fn prepare(&self, model: Self::Model) -> Result<Self::Runnable>;
283}
284
285pub trait RunnableInterface: Debug + Send + Sync {
286    type Tensor: TensorInterface;
287    type Fact: FactInterface;
288    type State: StateInterface<Tensor = Self::Tensor>;
289    fn run(&self, inputs: impl IntoInputs<Self::Tensor>) -> Result<Vec<Self::Tensor>> {
290        self.spawn_state()?.run(inputs.into_inputs()?)
291    }
292
293    fn input_count(&self) -> Result<usize>;
294    fn output_count(&self) -> Result<usize>;
295    fn input_fact(&self, id: usize) -> Result<Self::Fact>;
296
297    fn output_fact(&self, id: usize) -> Result<Self::Fact>;
298
299    fn input_facts(&self) -> Result<impl Iterator<Item = Self::Fact>> {
300        Ok((0..self.input_count()?)
301            .map(|ix| self.input_fact(ix))
302            .collect::<Result<Vec<_>>>()?
303            .into_iter())
304    }
305
306    fn output_facts(&self) -> Result<impl Iterator<Item = Self::Fact>> {
307        Ok((0..self.output_count()?)
308            .map(|ix| self.output_fact(ix))
309            .collect::<Result<Vec<_>>>()?
310            .into_iter())
311    }
312
313    fn property_keys(&self) -> Result<Vec<String>>;
314    fn property(&self, name: impl AsRef<str>) -> Result<Self::Tensor>;
315
316    fn spawn_state(&self) -> Result<Self::State>;
317
318    fn cost_json(&self) -> Result<String>;
319
320    fn profile_json<I, IV, IE>(&self, inputs: Option<I>) -> Result<String>
321    where
322        I: IntoIterator<Item = IV>,
323        IV: TryInto<Self::Tensor, Error = IE>,
324        IE: Into<anyhow::Error> + Debug;
325}
326
327pub trait StateInterface: Debug + Clone + Send {
328    type Fact: FactInterface;
329    type Tensor: TensorInterface;
330
331    fn input_count(&self) -> Result<usize>;
332    fn output_count(&self) -> Result<usize>;
333
334    fn run(&mut self, inputs: impl IntoInputs<Self::Tensor>) -> Result<Vec<Self::Tensor>>;
335}
336
337pub trait TensorInterface: Debug + Sized + Clone + PartialEq + Send + Sync {
338    fn datum_type(&self) -> Result<DatumType>;
339    fn from_bytes(dt: DatumType, shape: &[usize], data: &[u8]) -> Result<Self>;
340    fn as_bytes(&self) -> Result<(DatumType, &[usize], &[u8])>;
341
342    fn from_slice<T: Datum>(shape: &[usize], data: &[T]) -> Result<Self> {
343        let data = unsafe {
344            std::slice::from_raw_parts(data.as_ptr() as *const u8, std::mem::size_of_val(data))
345        };
346        Self::from_bytes(T::datum_type(), shape, data)
347    }
348
349    fn as_slice<T: Datum>(&self) -> Result<&[T]> {
350        let (dt, _shape, data) = self.as_bytes()?;
351        ensure!(T::datum_type() == dt);
352        let data = unsafe {
353            std::slice::from_raw_parts(
354                data.as_ptr() as *const T,
355                data.len() / std::mem::size_of::<T>(),
356            )
357        };
358        Ok(data)
359    }
360
361    fn as_shape_and_slice<T: Datum>(&self) -> Result<(&[usize], &[T])> {
362        let (_, shape, _) = self.as_bytes()?;
363        let data = self.as_slice()?;
364        Ok((shape, data))
365    }
366
367    fn shape(&self) -> Result<&[usize]> {
368        let (_, shape, _) = self.as_bytes()?;
369        Ok(shape)
370    }
371
372    fn convert_to(&self, to: DatumType) -> Result<Self>;
373}
374
375pub trait FactInterface: Debug + Display + Clone {
376    type Dim: DimInterface;
377    fn datum_type(&self) -> Result<DatumType>;
378    fn rank(&self) -> Result<usize>;
379    fn dim(&self, axis: usize) -> Result<Self::Dim>;
380
381    fn dims(&self) -> Result<impl Iterator<Item = Self::Dim>> {
382        Ok((0..self.rank()?).map(|axis| self.dim(axis)).collect::<Result<Vec<_>>>()?.into_iter())
383    }
384}
385
386pub trait DimInterface: Debug + Display + Clone {
387    fn eval(&self, values: impl IntoIterator<Item = (impl AsRef<str>, i64)>) -> Result<Self>;
388    fn to_int64(&self) -> Result<i64>;
389}
390
391pub trait InferenceFactInterface: Debug + Display + Default + Clone {
392    fn empty() -> Result<Self>;
393}
394
395pub trait AsFact<M, F>: Debug {
396    fn as_fact(&self, model: &M) -> Result<Bow<'_, F>>;
397}
398
399#[repr(C)]
400#[derive(Debug, PartialEq, Eq, Copy, Clone)]
401pub enum DatumType {
402    Bool = 0x01,
403    U8 = 0x11,
404    U16 = 0x12,
405    U32 = 0x14,
406    U64 = 0x18,
407    I8 = 0x21,
408    I16 = 0x22,
409    I32 = 0x24,
410    I64 = 0x28,
411    F16 = 0x32,
412    F32 = 0x34,
413    F64 = 0x38,
414    #[cfg(feature = "complex")]
415    ComplexI16 = 0x42,
416    #[cfg(feature = "complex")]
417    ComplexI32 = 0x44,
418    #[cfg(feature = "complex")]
419    ComplexI64 = 0x48,
420    #[cfg(feature = "complex")]
421    ComplexF16 = 0x52,
422    #[cfg(feature = "complex")]
423    ComplexF32 = 0x54,
424    #[cfg(feature = "complex")]
425    ComplexF64 = 0x58,
426}
427
428impl DatumType {
429    pub fn size_of(&self) -> usize {
430        use DatumType::*;
431        match &self {
432            Bool | U8 | I8 => 1,
433            U16 | I16 | F16 => 2,
434            U32 | I32 | F32 => 4,
435            U64 | I64 | F64 => 8,
436            #[cfg(feature = "complex")]
437            ComplexI16 | ComplexF16 => 4,
438            #[cfg(feature = "complex")]
439            ComplexI32 | ComplexF32 => 8,
440            #[cfg(feature = "complex")]
441            ComplexI64 | ComplexF64 => 16,
442        }
443    }
444
445    pub fn is_bool(&self) -> bool {
446        *self == DatumType::Bool
447    }
448
449    pub fn is_number(&self) -> bool {
450        *self != DatumType::Bool
451    }
452
453    pub fn is_unsigned(&self) -> bool {
454        use DatumType::*;
455        *self == U8 || *self == U16 || *self == U32 || *self == U64
456    }
457
458    pub fn is_signed(&self) -> bool {
459        use DatumType::*;
460        *self == I8 || *self == I16 || *self == I32 || *self == I64
461    }
462
463    pub fn is_float(&self) -> bool {
464        use DatumType::*;
465        *self == F16 || *self == F32 || *self == F64
466    }
467}
468
469pub trait Datum {
470    fn datum_type() -> DatumType;
471}
472
473// IntoInputs trait — ergonomic input conversion for run()
474pub trait IntoInputs<V: TensorInterface> {
475    fn into_inputs(self) -> Result<Vec<V>>;
476}
477
478// Arrays of anything convertible to Tensor
479impl<V, T, E, const N: usize> IntoInputs<V> for [T; N]
480where
481    V: TensorInterface,
482    T: TryInto<V, Error = E>,
483    E: Into<anyhow::Error>,
484{
485    fn into_inputs(self) -> Result<Vec<V>> {
486        self.into_iter().map(|v| v.try_into().map_err(|e| e.into())).collect()
487    }
488}
489
490// Vec<V> passthrough
491impl<V: TensorInterface> IntoInputs<V> for Vec<V> {
492    fn into_inputs(self) -> Result<Vec<V>> {
493        Ok(self)
494    }
495}
496
497// Tuples — each element converts independently
498macro_rules! impl_into_inputs_tuple {
499    ($($idx:tt : $T:ident),+) => {
500        impl<V, $($T),+> IntoInputs<V> for ($($T,)+)
501        where
502            V: TensorInterface,
503            $($T: TryInto<V>,
504              <$T as TryInto<V>>::Error: Into<anyhow::Error>,)+
505        {
506            fn into_inputs(self) -> Result<Vec<V>> {
507                Ok(vec![$(self.$idx.try_into().map_err(|e| e.into())?),+])
508            }
509        }
510    };
511}
512
513impl_into_inputs_tuple!(0: A);
514impl_into_inputs_tuple!(0: A, 1: B);
515impl_into_inputs_tuple!(0: A, 1: B, 2: C);
516impl_into_inputs_tuple!(0: A, 1: B, 2: C, 3: D);
517impl_into_inputs_tuple!(0: A, 1: B, 2: C, 3: D, 4: E_);
518impl_into_inputs_tuple!(0: A, 1: B, 2: C, 3: D, 4: E_, 5: F);
519impl_into_inputs_tuple!(0: A, 1: B, 2: C, 3: D, 4: E_, 5: F, 6: G);
520impl_into_inputs_tuple!(0: A, 1: B, 2: C, 3: D, 4: E_, 5: F, 6: G, 7: H);
521
522/// Convert any compatible input into a `V: TensorInterface`.
523pub fn tensor<V, T, E>(v: T) -> Result<V>
524where
525    V: TensorInterface,
526    T: TryInto<V, Error = E>,
527    E: Into<anyhow::Error>,
528{
529    v.try_into().map_err(|e| e.into())
530}
531
532macro_rules! impl_datum_type {
533    ($ty:ty, $c_repr:expr) => {
534        impl Datum for $ty {
535            fn datum_type() -> DatumType {
536                $c_repr
537            }
538        }
539    };
540}
541
542impl_datum_type!(bool, DatumType::Bool);
543impl_datum_type!(u8, DatumType::U8);
544impl_datum_type!(u16, DatumType::U16);
545impl_datum_type!(u32, DatumType::U32);
546impl_datum_type!(u64, DatumType::U64);
547impl_datum_type!(i8, DatumType::I8);
548impl_datum_type!(i16, DatumType::I16);
549impl_datum_type!(i32, DatumType::I32);
550impl_datum_type!(i64, DatumType::I64);
551impl_datum_type!(half::f16, DatumType::F16);
552impl_datum_type!(f32, DatumType::F32);
553impl_datum_type!(f64, DatumType::F64);
tract_api/lib.rs

tract_api/
lib.rs
