tract_pulse/
lib.rs

1#![allow(clippy::len_zero)]
2#[macro_use]
3pub mod macros;
4
5pub mod fact;
6pub mod model;
7pub mod ops;
8
9pub mod internal {
10    pub use std::fmt;
11    pub use tract_nnef::internal::*;
12    pub use tract_pulse_opl::tract_nnef;
13
14    pub use downcast_rs::Downcast;
15
16    pub use crate::fact::PulsedFact;
17    pub use crate::model::{PulsedModel, PulsedModelExt};
18    pub use crate::ops::{OpPulsifier, PulsedOp};
19}
20
21use std::ops::ControlFlow;
22
23use internal::*;
24
25pub use ops::PulsedOp;
26
27pub trait WithPulse {
28    fn enable_pulse(&mut self);
29    fn with_pulse(self) -> Self;
30}
31
32impl WithPulse for tract_nnef::framework::Nnef {
33    fn enable_pulse(&mut self) {
34        self.enable_tract_core();
35        self.registries.push(tract_nnef_registry());
36    }
37    fn with_pulse(mut self) -> Self {
38        self.enable_pulse();
39        self
40    }
41}
42
43pub fn tract_nnef_registry() -> Registry {
44    let mut reg = tract_pulse_opl::tract_nnef_registry();
45    ops::delay::register(&mut reg);
46    reg.extensions.push(Box::new(decl_stream_symbol));
47    reg
48}
49
50fn decl_stream_symbol(
51    _proto_model: &mut ModelBuilder,
52    name: &Identifier,
53    _rest: &str,
54) -> TractResult<ControlFlow<(), ()>> {
55    if name.0 == "tract_pulse_streaming_symbol" {
56        Ok(ControlFlow::Break(()))
57    } else {
58        Ok(ControlFlow::Continue(()))
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn test_source_must_stream() {
68        let mut model = TypedModel::default();
69        let s = model.symbols.sym("S");
70        let _a = model.add_source("a", f32::fact([1, 2, 3])).unwrap();
71        model.auto_outputs().unwrap();
72        assert!(PulsedModel::new(&model, s.clone(), &4.to_dim()).is_err());
73
74        let mut model = TypedModel::default();
75        let _a = model.add_source("a", f32::fact(dims![1, s, 3].as_ref())).unwrap();
76        model.auto_outputs().unwrap();
77        let pulse = PulsedModel::new(&model, s, &4.to_dim()).unwrap();
78        assert_eq!(
79            *pulse.outlet_fact(OutletId::new(0, 0)).unwrap().to_typed_fact().unwrap(),
80            f32::fact([1usize, 4, 3])
81        );
82    }
83
84    #[test]
85    fn test_immediate() {
86        let mut model = TypedModel::default();
87        let s = model.symbols.sym("S");
88        let _a = model.add_source("a", f32::fact(dims![s, 2, 3].as_ref())).unwrap();
89        model.auto_outputs().unwrap();
90
91        let pulse = PulsedModel::new(&model, s, &4.to_dim()).unwrap();
92
93        assert_eq!(*pulse.input_fact(0).unwrap().to_typed_fact().unwrap(), f32::fact([4, 2, 3]));
94        assert_eq!(*pulse.output_fact(0).unwrap().to_typed_fact().unwrap(), f32::fact([4, 2, 3]));
95    }
96}
tract_pulse/lib.rs

tract_pulse/
lib.rs
