Pulse v2

cli/src/main.rs

@@ -74,6 +74,9 @@ pub const STAGES: &[&str] = &[
     "pulse",
     "pulse-to-type",
     "pulse-declutter",
+    "pulse-v2",
+    "pulse-v2-to-type",
+    "pulse-v2-declutter",
     "set",
     "set-declutter",
     "nnef-cycle",
@@ -169,6 +172,7 @@ fn main() -> TractResult<()> {
         .arg(arg!(-O --optimize "Optimize before running"))
         .arg(arg!(--"assert-maximal-mm-quality-cost" [MAX] "Maximum value for quality category (0=assembly, 4=dreadful rust code)"))
         .arg(arg!(--pulse [PULSE] "Translate to pulse network"))
+        .arg(arg!(--"pulse-v2" [SYM] "Translate to pulse-v2 network (streaming axis symbol, default S)"))
 
         .arg(arg!(--"machine-friendly" "Machine friendly output"))
         .arg(arg!(--"timeout" [SECONDS] "Kill the process after this many seconds"))

cli/src/params.rs

@@ -725,6 +725,16 @@ impl Parameters {
                 stage!("pulse-to-type", pulsed_model -> typed_model, |m:PulsedModel| m.into_typed());
                 stage!("pulse-declutter", typed_model -> typed_model, |m:TypedModel| m.into_decluttered());
             }
+            if let Some(spec) = matches.get_one::<String>("pulse-v2") {
+                stage!("pulse-v2", typed_model -> typed_model, |m:TypedModel| {
+                    use tract_pulse::v2::PulseV2Model;
+                    let stream_sym = m.symbols.sym(spec);
+                    let pv2 = PulseV2Model::new(&m, stream_sym)?;
+                    pv2.into_typed()
+                });
+                stage!("pulse-v2-to-type", typed_model -> typed_model, |m:TypedModel| Ok(m));
+                stage!("pulse-v2-declutter", typed_model -> typed_model, |m:TypedModel| m.into_decluttered());
+            }
         }
         let mut transforms: Vec<&str> = matches
             .get_many::<String>("transform")

data/src/dim/tree.rs

@@ -1902,3 +1902,17 @@ mod tests {
         assert_eq!(a_s, c_s, "8*(-1*B) should simplify the same as -8*B");
     }
 }
+
+#[test]
+fn mul_neg_b_by_8() {
+    let s = crate::dim::SymbolScope::default();
+    let b = Sym(s.sym("B"));
+    // 8*(-1*B) should equal -8*B
+    let a = Mul(vec![Val(8), MulInt(-1, Box::new(b.clone()))]);
+    let c = MulInt(-8, Box::new(b.clone()));
+    let a_s = a.simplify();
+    let c_s = c.simplify();
+    eprintln!("8*(-1*B) = {a_s}");
+    eprintln!("-8*B     = {c_s}");
+    assert_eq!(a_s, c_s, "8*(-1*B) should simplify the same as -8*B");
+}

harness/core-proptest-pulse/Cargo.toml

@@ -13,3 +13,4 @@ tract-pulse.workspace = true
 log.workspace = true
 proptest.workspace = true
 env_logger.workspace = true
+tract-tensorflow.workspace = true

harness/core-proptest-pulse/src/conv_plus_conv.rs

@@ -63,7 +63,7 @@ impl Arbitrary for ConvOp {
 }
 
 #[derive(Debug, Clone)]
-struct ConvPlusConvProblem {
+pub(crate) struct ConvPlusConvProblem {
     input: Tensor,
     pulse: usize,
     convs: Vec<ConvOp>,
@@ -123,6 +123,10 @@ impl ConvPlusConvProblem {
             2,
         )
     }
+
+    pub fn run_v2(&self) -> TestCaseResult {
+        crate::v2::run_and_compare_v2(Self::model(&self.convs), self.pulse, &self.input, 2)
+    }
 }
 
 proptest! {

harness/core-proptest-pulse/src/deconv.rs

@@ -68,7 +68,7 @@ impl Arbitrary for DeconvOp {
 }
 
 #[derive(Debug, Clone)]
-struct DeconvProblem {
+pub(crate) struct DeconvProblem {
     input: Array3<f32>,
     pulse: usize,
     deconv: DeconvOp,
@@ -94,6 +94,17 @@ impl Arbitrary for DeconvProblem {
 }
 
 impl DeconvProblem {
+    pub fn run_v2(&self) -> TestCaseResult {
+        let mut model = TypedModel::default();
+        let mut fact = f32::fact(self.input.shape());
+        let s = model.symbols.sym("S");
+        fact.shape.set(2, s.to_dim());
+        let input = model.add_source("a", fact).unwrap();
+        let id = self.deconv.chain("deconv1", &mut model, input);
+        model.select_output_outlets(&[id]).unwrap();
+        crate::v2::run_and_compare_v2(model, self.pulse, &self.input.clone().into_tensor(), 2)
+    }
+
     pub fn run(&self) -> TestCaseResult {
         let mut model = TypedModel::default();
         let mut fact = f32::fact(self.input.shape());

harness/core-proptest-pulse/src/delay_plus_downsample.rs

@@ -6,7 +6,7 @@ use tract_core::tract_data::itertools::Itertools;
 use super::*;
 
 #[derive(Debug, Clone)]
-struct DelayPlusDownsampleProblem {
+pub(crate) struct DelayPlusDownsampleProblem {
     input: usize,
     pulse: usize,
     delay: usize,
@@ -45,6 +45,22 @@ impl Arbitrary for DelayPlusDownsampleProblem {
 }
 
 impl DelayPlusDownsampleProblem {
+    pub fn run_v2(&self) -> TestCaseResult {
+        let mut model = TypedModel::default();
+        let s = model.symbols.sym("S");
+        let a = model.add_source("a", f32::fact(dims!(1, s, 1))).unwrap();
+        let crop = model.wire_node("delay", Slice::new(1, self.delay, s), &[a]).unwrap();
+        let ds = model
+            .wire_node(
+                "ds",
+                Downsample { axis: 1, stride: self.stride as isize, modulo: self.modulo },
+                &crop,
+            )
+            .unwrap();
+        model.select_output_outlets(&ds).unwrap();
+        crate::v2::run_and_compare_v2(model, self.pulse, &t(self.input).into_tensor(), 1)
+    }
+
     pub fn run(&self) -> TestCaseResult {
         let mut model = TypedModel::default();
         let s = model.symbols.sym("S");

harness/core-proptest-pulse/src/delay_plus_pool.rs

@@ -5,7 +5,7 @@ use tract_core::ops::cnn::MaxPool;
 use super::*;
 
 #[derive(Debug, Clone)]
-struct DelayPlusPoolProblem {
+pub(crate) struct DelayPlusPoolProblem {
     input: Vec<f32>,
     pulse: usize,
     delay: usize,
@@ -44,6 +44,27 @@ impl Arbitrary for DelayPlusPoolProblem {
 }
 
 impl DelayPlusPoolProblem {
+    pub fn run_v2(&self) -> TestCaseResult {
+        let mut model = TypedModel::default();
+        let s = model.symbols.sym("S");
+        let a = model.add_source("a", f32::fact(dims!(1, s, 1))).unwrap();
+        let crop = model.wire_node("delay", Slice::new(1, self.delay, s), &[a]).unwrap();
+        let pool_spec = PoolSpec::new(
+            DataFormat::NHWC,
+            tvec!(self.pool_window),
+            self.padding.clone(),
+            None,
+            Some(tvec!(self.stride)),
+            1,
+            1,
+        );
+        let pool = model.wire_node("pool", MaxPool::new(pool_spec, None), &crop).unwrap();
+        model.select_output_outlets(&pool).unwrap();
+        let input =
+            arr1(&self.input).into_shape_with_order((1, self.input.len(), 1)).unwrap().into_dyn();
+        crate::v2::run_and_compare_v2(model, self.pulse, &input.into_tensor(), 1)
+    }
+
     pub fn run(&self) -> TestCaseResult {
         let mut model = TypedModel::default();
         let s = model.symbols.sym("S");

harness/core-proptest-pulse/src/lib.rs

@@ -20,12 +20,13 @@ use tract_core::ops::nn::DataFormat;
 use tract_ndarray::prelude::*;
 use tract_pulse::internal::*;
 
-mod conv_plus_conv;
-mod deconv;
-mod delay_plus_downsample;
-mod delay_plus_pool;
+pub(crate) mod conv_plus_conv;
+pub(crate) mod deconv;
+pub(crate) mod delay_plus_downsample;
+pub(crate) mod delay_plus_pool;
 mod einsum;
-mod pad_plus_conv;
+pub(crate) mod pad_plus_conv;
+mod v2;
 
 #[allow(dead_code)]
 fn setup_test_logger() {

harness/core-proptest-pulse/src/pad_plus_conv.rs

@@ -5,7 +5,7 @@ use proptest::*;
 use super::*;
 
 #[derive(Debug, Clone)]
-struct PadPlusConvProblem {
+pub(crate) struct PadPlusConvProblem {
     pad_before: usize,
     pad_after: usize,
     pad_mode: PadMode,
@@ -107,6 +107,50 @@ proptest! {
     fn proptest_conv(pb in PadPlusConvProblem::arbitrary()) { pb.run().unwrap() }
 }
 
+impl PadPlusConvProblem {
+    pub fn run_v2(&self) -> TestCaseResult {
+        let mut model = TypedModel::default();
+        let s = model.symbols.sym("S");
+        let mut wire = model.add_source("a", f32::fact(dims!(1, 1, s))).unwrap();
+        if self.pad_before > 0 || self.pad_after > 0 {
+            wire = model
+                .wire_node(
+                    "pad",
+                    Pad::new(
+                        vec![(0, 0), (0, 0), (self.pad_before, self.pad_after)],
+                        self.pad_mode.clone(),
+                    ),
+                    &[wire],
+                )
+                .unwrap()[0];
+        }
+        let kernel = model.add_const("kernel", self.ker.clone()).unwrap();
+        let bias = model.add_const("bias", tensor0(0f32)).unwrap();
+        let conv = model
+            .wire_node(
+                "conv",
+                Conv {
+                    pool_spec: PoolSpec {
+                        data_format: DataFormat::NCHW,
+                        kernel_shape: self.ker.shape()[2..].into(),
+                        padding: PaddingSpec::Valid,
+                        dilations: Some(tvec!(self.dilation)),
+                        strides: Some(tvec!(self.stride)),
+                        input_channels: 1,
+                        output_channels: 1,
+                    },
+                    kernel_fmt: tract_core::ops::cnn::KernelFormat::OIHW,
+                    group: 1,
+                    q_params: None,
+                },
+                &[wire, kernel, bias],
+            )
+            .unwrap();
+        model.select_output_outlets(&conv).unwrap();
+        crate::v2::run_and_compare_v2(model, self.pulse, &self.input.clone().into_tensor(), 2)
+    }
+}
+
 #[test]
 fn conv_1() {
     PadPlusConvProblem {