Skip to main content

tract_hir/infer/
mod.rs

1use crate::internal::*;
2
3#[macro_use]
4pub mod helpers;
5#[macro_use]
6pub mod rules;
7
8mod analyser;
9mod fact;
10mod factoid;
11mod model;
12mod ops;
13mod optim;
14
15pub use self::fact::InferenceFact;
16pub use self::factoid::*;
17pub use self::model::InferenceModelExt;
18pub use self::ops::InferenceOp;
19pub use self::rules::InferenceResult;
20pub use self::rules::InferenceRulesOp;
21pub use self::rules::Solver;
22pub use self::rules::TensorProxy;
23pub use self::rules::expr::IntoExp;
24pub use self::rules::expr::ToDimExp;
25
26pub fn check_input_arity(inputs: &[TensorProxy], expected: usize) -> TractResult<()> {
27    if inputs.len() != expected {
28        bail!("Wrong input number. Rules expect {}, node has {}.", expected, inputs.len())
29    } else {
30        Ok(())
31    }
32}
33
34pub fn check_output_arity(outputs: &[TensorProxy], expected: usize) -> TractResult<()> {
35    if outputs.len() != expected {
36        bail!("Wrong output number. Rules expect {}, node has {}.", expected, outputs.len())
37    } else {
38        Ok(())
39    }
40}
41
42/// A model with partially types and shapes, as produced by parsing ONNX or
43/// Tensorflow graphs.
44pub type InferenceModel = Graph<InferenceFact, Box<dyn InferenceOp>>;
45/// Node for InferenceModel graph
46pub type InferenceNode = Node<InferenceFact, Box<dyn InferenceOp>>;
47/// A ModelPatch for InferenceModel.
48pub type InferenceModelPatch = ModelPatch<InferenceFact, Box<dyn InferenceOp>>;
49/// An execution plan for InferenceModel.
50pub type InferenceSimplePlan = SimplePlan<InferenceFact, Box<dyn InferenceOp>>;
51/// An execution state for InferenceModel.
52pub type InferenceSimpleState = SimpleState<InferenceFact, Box<dyn InferenceOp>>;
53
54impl<'a> From<&'a Box<dyn InferenceOp>> for Box<dyn InferenceOp> {
55    fn from(it: &'a Box<dyn InferenceOp>) -> Box<dyn InferenceOp> {
56        tract_core::dyn_clone::clone_box(it.as_ref())
57    }
58}