Repository: webonnx/wonnx (language: python)

Issue:
fix: remove logic that removes identity ops for now, it fails in certain scenarios
Temporary fix for #43 (see issue for more explanation)


Discussion / hints:

Fix patch:
diff --git a/Cargo.toml b/Cargo.toml
index 1a9c595e..648d1515 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -23,7 +23,7 @@ exclude = [
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-wgpu= "0.12.0"
+wgpu = "0.12.0"
 bytemuck = "1.7.2"
 env_logger = "0.9.0"
 pollster = "0.2.4"
diff --git a/README.md b/README.md
index 6a513f19..67f780c1 100644
--- a/README.md
+++ b/README.md
@@ -51,6 +51,15 @@ inputs = {"x": [-1.0, 2.0]}
 assert session.run(inputs) == {"y": [0.0, 2.0]}
 ```
 
+To build the Python module for development:
+
+````sh
+python3 -m venv .env
+source .env/bin/activate
+pip install maturin
+maturin develop
+````
+
 ## Running a model from scratch
 
 - To run an onnx model, first simplify it with [onnx-simplifier](https://github.com/daquexian/onnx-simplifier), with the command:
@@ -116,7 +125,6 @@ To implement an operator all you have to do is:
 6. Write the logic.
 
 There is default variables in the context: 
-- `{{ inputs[0] }}`: the name of the variable.
 - `{{ i_lens[0] }}`: the length of the input 0. This also work for output: `{{ o_lens[0] }}` and other input `{{ i_lens[1] }}`
 - `{{ i_shape[0] }}`: the array of dimensions of input 0. To get the first dimension of the array, just use: `{{ i_shape[0][0] }}` 
 - `{{ i_chunks[0] }}`: the size of the chunks of each dimensions of input 0. By default, each variable is represented as a long array of values where to get to specific values you have to move by chunks. Those chunks are represented within this variable. To get the size of the chunks of the first dimensions use: `{{ i_chunks[0][0] }}`.
diff --git a/examples/mnist.rs b/examples/mnist.rs
index 70bbf635..a0cf82a8 100644
--- a/examples/mnist.rs
+++ b/examples/mnist.rs
@@ -30,7 +30,7 @@ async fn execute_gpu() -> Option<HashMap<String, Vec<f32>>> {
         .await
         .unwrap();
     let time_pre_compute = Instant::now();
-    let result = session.run(input_data).await.unwrap();
+    let result = session.run(&input_data).await.unwrap();
     let time_post_compute = Instant::now();
     println!(
         "time: post_compute: {:#?}",
diff --git a/examples/simple_graph.rs b/examples/simple_graph.rs
index ad7597a4..a6a873a9 100644
--- a/examples/simple_graph.rs
+++ b/examples/simple_graph.rs
@@ -49,7 +49,7 @@ async fn execute_gpu() -> Result<HashMap<String, Vec<f32>>, SessionError> {
         .await
         .expect("Session did not create");
 
-    session.run(input_data).await
+    session.run(&input_data).await
 }
 
 // #[wasm_bindgen_test]
diff --git a/examples/squeeze.rs b/examples/squeeze.rs
index acfbf105..f54e7b49 100644
--- a/examples/squeeze.rs
+++ b/examples/squeeze.rs
@@ -16,20 +16,13 @@ use wonnx::WonnxError;
 // Args Management
 async fn run() {
     let probabilities = &execute_gpu().await.unwrap();
-
     let (_, probabilities) = probabilities.iter().next().unwrap();
-
     let mut probabilities = probabilities.iter().enumerate().collect::<Vec<_>>();
-
     probabilities.sort_unstable_by(|a, b| b.1.partial_cmp(a.1).unwrap());
 
     let class_labels = get_imagenet_labels();
 
     for i in 0..10 {
-        //      println!(
-        //          "probabilities[i*9..(i+1)*9]: {:#?}",
-        //          &probabilities[i * 12..(i + 1) * 12]
-        //      );
         println!(
             "Infered result: {} of class: {}",
             class_labels[probabilities[i].0], probabilities[i].0
@@ -46,7 +39,7 @@ async fn execute_gpu() -> Result<HashMap<String, Vec<f32>>, WonnxError> {
     let session = wonnx::Session::from_path("examples/data/models/opt-squeeze.onnx").await?;
     let time_pre_compute = Instant::now();
     info!("Start Compute");
-    let result = session.run(input_data.clone()).await?;
+    let result = session.run(&input_data).await?;
     let time_post_compute = Instant::now();
     println!(
         "time: first_prediction: {:#?}",
diff --git a/examples/yolo.rs b/examples/yolo.rs
deleted file mode 100644
index e2f6eeb3..00000000
--- a/examples/yolo.rs
+++ /dev/null
@@ -1,50 +0,0 @@
-use std::collections::HashMap;
-// Indicates a f32 overflow in an intermediate Collatz value
-// use wasm_bindgen_test::*;
-
-use std::time::Instant;
-// Args Management
-async fn run() {
-    let _result = execute_gpu().await;
-    //  println!("steps: {:#?}", &result);
-    // println!("steps[1..5]: {:#?}", &steps[0..5]);
-    #[cfg(target_arch = "wasm32")]
-    // log::info!("steps[0..5]: {:#?}", &steps[0..5]);
-    assert_eq!(steps[0..5], [0.0, 0.0, 0.0, 0.0, 0.0]);
-}
-
-// Hardware management
-async fn execute_gpu() -> HashMap<String, Vec<f32>> {
-    let n = 416;
-    let c = 3;
-    let mut input_data = HashMap::new();
-
-    let data: Vec<f32> = (0..n * n * c).map(|x| x as f32).collect();
-    input_data.insert("image".to_string(), data.as_slice());
-
-    let session = wonnx::Session::from_path("examples/data/models/opt-tinyyolov2-8.onnx")
-        .await
-        .unwrap();
-    let time_pre_compute = Instant::now();
-    let result = session.run(input_data).await.unwrap();
-    let time_post_compute = Instant::now();
-    println!(
-        "time: post_compute: {:#?}",
-        time_post_compute - time_pre_compute
-    );
-    result
-}
-
-fn main() {
-    #[cfg(not(target_arch = "wasm32"))]
-    {
-        env_logger::init();
-        pollster::block_on(run());
-    }
-    #[cfg(target_arch = "wasm32")]
-    {
-        // std::panic::set_hook(Box::new(console_error_panic_hook::hook));
-        //  console_log::init().expect("could not initialize logger");
-        wasm_bindgen_futures::spawn_local(run());
-    }
-}
diff --git a/py-wonnx/src/lib.rs b/py-wonnx/src/lib.rs
index 06f25638..e1e1a4d9 100644
--- a/py-wonnx/src/lib.rs
+++ b/py-wonnx/src/lib.rs
@@ -29,7 +29,7 @@ impl PySession {
         for (key, value) in map.iter() {
             inputs.insert(key.clone(), value.as_slice());
         }
-        let result = pollster::block_on(self.session.run(inputs)).unwrap();
+        let result = pollster::block_on(self.session.run(&inputs)).unwrap();
 
         Ok(result)
     }
diff --git a/src/compiler.rs b/src/compiler.rs
index 2a3bc2d9..4872c51c 100644
--- a/src/compiler.rs
+++ b/src/compiler.rs
@@ -1,14 +1,15 @@
 use crate::utils::{ceil, get_attribute, AttributeNotFoundError, DataType, Shape};
-use std::collections::HashMap;
 use tera::{Context, Tera};
 use thiserror::Error;
 
-// Escaping special characters as well as adding `var_` in the beginning of the variable name to avoid collisions with wgsl syntax.
-// FIXME: this has the potential to cause collisions (i.e. "a/b" and "a.b" will both translate to "ab" and hence cannot be used simultaneously)
-fn to_wgsl_variable_name(input_or_output_name: &str) -> String {
-    String::from("var_")
-        + &input_or_output_name.replace(&['(', ')', ',', '\"', '.', ';', ':', '\'', '/'][..], "")
-}
+/// The maximum number of threads that can be spawned in each dimension, according to the WebGPU specification. See
+// https://www.w3.org/TR/webgpu/#dom-supported-limits-maxcomputeworkgroupsperdimension
+pub const MAX_COMPUTE_WORKGROUPS_PER_DIMENSION: u32 = 65535;
+
+/// The maximum workgroup size per dimension (see https://www.w3.org/TR/webgpu/#dom-supported-limits-maxcomputeworkgroupsizex)
+pub const MAX_WORKGROUP_SIZE_X: u32 = 256;
+pub const MAX_WORKGROUP_SIZE_Y: u32 = 256;
+pub const MAX_WORKGROUP_SIZE_Z: u32 = 64;
 
 lazy_static! {
     // Templates for shader source code that we generate for nodes
@@ -29,11 +30,6 @@ lazy_static! {
             include_str!("../templates/endomorphism/batchnormalization.wgsl"),
         )
         .unwrap();
-        tera.add_raw_template(
-            "endomorphism/copy.wgsl",
-            include_str!("../templates/endomorphism/copy.wgsl"),
-        )
-        .unwrap();
         tera.add_raw_template(
             "endomorphism/softmax.wgsl",
             include_str!("../templates/endomorphism/softmax.wgsl"),
@@ -147,91 +143,57 @@ pub enum CompileError {
         input_index: usize,
         input_shape: Shape,
     },
+
+    #[error("the model exceeds the limit for {0}: {1} > {2}")]
+    ComputeLimitExceeded(String, u32, u32),
 }
 
 pub fn compile(
     node: &crate::onnx::NodeProto,
-    shape_infos: &HashMap<String, Shape>,
+    input_shapes: &[&Shape],
+    output_shapes: &[&Shape],
     opset_version: i64,
 ) -> Result<CompiledNode, CompileError> {
-    // Escape unwanted characters
-    let mut inputs = node.get_input().to_vec();
-    let mut outputs = node.get_output().to_vec();
-
-    let input_shape = inputs
-        .iter()
-        .map(|input| match shape_infos.get(input.as_str()) {
-            Some(info) => Ok(info),
-            None => Err(CompileError::DimensionsMissing(
-                input.to_string(),
-                node.get_name().to_string(),
-            )),
-        })
-        .collect::<Result<Vec<_>, CompileError>>()?;
-
-    let output_shape = outputs
-        .iter()
-        .map(|output| match shape_infos.get(output.as_str()) {
-            Some(info) => Ok(info),
-            None => Err(CompileError::DimensionsMissing(
-                output.to_string(),
-                node.get_name().to_string(),
-            )),
-        })
-        .collect::<Result<Vec<_>, CompileError>>()?;
-
-    let input_lengths = input_shape
+    let input_lengths = input_shapes
         .iter()
         .map(|shape| shape.element_count())
         .collect::<Vec<u64>>();
 
-    let output_lengths = output_shape
+    let output_lengths = output_shapes
         .iter()
         .map(|shape| shape.element_count())
         .collect::<Vec<u64>>();
 
-    // Generate variable names from the input names (which may contain special characters we don't want)
-    inputs = inputs
-        .iter()
-        .map(|input| to_wgsl_variable_name(input))
-        .collect::<Vec<_>>();
-
-    outputs = outputs
-        .iter()
-        .map(|output| to_wgsl_variable_name(output))
-        .collect::<Vec<_>>();
-
-    let input_chunks: Vec<Vec<u64>> = input_shape.iter().map(|d| d.chunks()).collect();
-    let output_chunks: Vec<Vec<u64>> = output_shape.iter().map(|d| d.chunks()).collect();
+    let input_chunks: Vec<Vec<u64>> = input_shapes.iter().map(|d| d.chunks()).collect();
+    let output_chunks: Vec<Vec<u64>> = output_shapes.iter().map(|d| d.chunks()).collect();
 
     let mut context = Context::new();
-    context.insert("inputs", &inputs);
-    context.insert("outputs", &outputs);
     context.insert("i_lens", &input_lengths);
     context.insert("o_lens", &output_lengths);
-    context.insert("i_shape", &input_shape);
-    context.insert("o_shape", &output_shape);
+    context.insert("i_shape", &input_shapes);
+    context.insert("o_shape", &output_shapes);
     context.insert("i_chunks", &input_chunks);
     context.insert("o_chunks", &output_chunks);
     context.insert("op_type", &node.get_op_type());
     context.insert("opset_version", &opset_version);
 
     let (template, x, y, z) = match node.get_op_type() {
+        op @ ("Reshape" | "Dropout" | "Identity" | "Flatten" | "Squeeze" | "Unsqueeze") => {
+            // These ops should all be optimized away earlier
+            return Err(CompileError::InvalidOperation(op.to_string()));
+        }
+
         // Map simple function
         "Abs" | "Acos" | "Asin" | "Atan" | "Ceil" | "Cos" | "Cosh" | "Exp" | "Floor" | "Log"
-        | "Round" | "Sign" | "Sin" | "Sinh" | "Sqrt" | "Tan" | "Tanh" => (
-            "endomorphism/map.wgsl".to_string(),
-            ceil(output_lengths[0], 4) as _,
-            1,
-            1,
-        ),
-        // Copy data
-        "Reshape" | "Dropout" | "Flatten" | "Squeeze" | "Unsqueeze" | "Identity" => (
-            "endomorphism/copy.wgsl".to_string(),
-            ceil(output_lengths[0], 16) as _,
-            1,
-            1,
-        ),
+        | "Round" | "Sign" | "Sin" | "Sinh" | "Sqrt" | "Tan" | "Tanh" => {
+            let (x_threads, workgroup_size_x) = workgroup_size(
+                ceil(output_lengths[0], 4),
+                MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
+                MAX_WORKGROUP_SIZE_X,
+            )?;
+            context.insert("workgroup_size_x", &workgroup_size_x);
+            ("endomorphism/map.wgsl".to_string(), x_threads, 1, 1)
+        }
 
         "Softmax" => {
             let default_axis = match opset_version {
@@ -247,7 +209,7 @@ pub fn compile(
             let mut axis = get_attribute("axis", Some(default_axis), node)?;
             if axis < 0 {
                 if opset_version >= 13 {
-                    axis += input_shape[0].rank() as i64;
+                    axis += input_shapes[0].rank() as i64;
                 } else {
                     return Err(CompileError::InvalidAttributeValue {
                         attribute: "axis".to_string(),
@@ -257,7 +219,7 @@ pub fn compile(
                 }
             }
 
-            if axis >= (input_shape[0].rank() as i64) {
+            if axis >= (input_shapes[0].rank() as i64) {
                 return Err(CompileError::InvalidAttributeValue {
                     attribute: "axis".to_string(),
                     value: format!("{}", axis),
@@ -305,12 +267,15 @@ pub fn compile(
                     }
                 },
             );
-            (
-                "endomorphism/arithmetic.wgsl".to_string(),
+
+            let (x_threads, workgroup_size_x) = workgroup_size(
                 ceil(output_lengths[0], 1024) as _,
-                1,
-                1,
-            )
+                MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
+                MAX_WORKGROUP_SIZE_X,
+            )?;
+            context.insert("workgroup_size_x", &workgroup_size_x);
+
+            ("endomorphism/arithmetic.wgsl".to_string(), x_threads, 1, 1)
         }
         // Not taking into account attributes
         "BatchNormalization" => {
@@ -334,26 +299,26 @@ pub fn compile(
 
             // [N,C,w,h] => [N,C,w,h] where [w,h] is normalized using stats for each [N,C]
             // N and C are optional and assumed to be one for lower-rank inputs
-            if input_shape[0].rank() <= 2 || input_shape[0].rank() > 4 {
+            if input_shapes[0].rank() <= 2 || input_shapes[0].rank() > 4 {
                 return Err(CompileError::UnimplementedVariant {
                     op: "BatchNormalization".to_string(),
-                    variant: format!("with input {}", input_shape[0]),
+                    variant: format!("with input {}", input_shapes[0]),
                 });
             }
 
-            let (input_batches, input_channels, input_w, input_h) = match input_shape[0].rank() {
-                2 => (1, 1, input_shape[0].dim(0), input_shape[0].dim(1)), // WxH, C=1, N=1
+            let (input_batches, input_channels, input_w, input_h) = match input_shapes[0].rank() {
+                2 => (1, 1, input_shapes[0].dim(0), input_shapes[0].dim(1)), // WxH, C=1, N=1
                 3 => (
                     1,
-                    input_shape[0].dim(0),
-                    input_shape[0].dim(1),
-                    input_shape[0].dim(2),
+                    input_shapes[0].dim(0),
+                    input_shapes[0].dim(1),
+                    input_shapes[0].dim(2),
                 ), // CxWxH, single batch N=1
                 4 => (
-                    input_shape[0].dim(0),
-                    input_shape[0].dim(1),
-                    input_shape[0].dim(2),
-                    input_shape[0].dim(3),
+                    input_shapes[0].dim(0),
+                    input_shapes[0].dim(1),
+                    input_shapes[0].dim(2),
+                    input_shapes[0].dim(3),
                 ), // NxCxWxH
                 _ => unreachable!(),
             };
@@ -361,7 +326,7 @@ pub fn compile(
             if input_batches == 0 || input_channels == 0 {
                 return Err(CompileError::InvalidInputShape {
                     input_index: 0,
-                    input_shape: input_shape[0].clone(),
+                    input_shape: input_shapes[0].clone(),
                 });
             }
 
@@ -396,9 +361,18 @@ pub fn compile(
         "Relu" | "Sigmoid" | "Softsign" | "Softplus" | "Clip" | "Celu" | "Elu" | "LeakyRelu" => {
             let alpha = get_attribute("alpha", Some(1.0), node)?;
             context.insert("alpha", &alpha);
+
+            let (x_threads, workgroup_size_x) = workgroup_size(
+                ceil(output_lengths[0], 4),
+                MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
+                MAX_WORKGROUP_SIZE_X,
+            )?;
+
+            context.insert("workgroup_size_x", &workgroup_size_x);
+
             (
                 "endomorphism/activation.wgsl".to_string(),
-                ceil(output_lengths[0], 4) as _,
+                x_threads as _,
                 1,
                 1,
             )
@@ -421,7 +395,7 @@ pub fn compile(
         op @ ("MaxPool" | "AveragePool" | "Conv" | "ConvRelu" | "ConvLeakyRelu" | "ConvMish"
         | "GlobalAveragePool") => {
             // TODO: Conv only support NxCxHxW for the moment.
-            debug_assert!(input_shape[0].rank() == 4);
+            debug_assert!(input_shapes[0].rank() == 4);
 
             // GlobalAveragePool is equivalent to AveragePool, with the kernel shape set to the size of the input tensor
             // See https://github.com/onnx/onnx/blob/main/docs/Operators.md#globalaveragepool
@@ -435,7 +409,7 @@ pub fn compile(
             let auto_pad = get_attribute("auto_pad", Some("NOTSET".to_string()), node)?;
             let dilations = get_attribute("dilations", Some(vec![1, 1]), node)?;
             let kernel_shape = if is_global_average_pool {
-                vec![input_shape[0].dim(2) as i64, input_shape[0].dim(3) as i64]
+                vec![input_shapes[0].dim(2) as i64, input_shapes[0].dim(3) as i64]
             } else {
                 get_attribute::<Vec<i64>>("kernel_shape", None, node)?
             };
@@ -480,8 +454,8 @@ pub fn compile(
                 }
             };
 
-            let input_shape = input_shape[0];
-            let output_shape = output_shape[0];
+            let input_shape = &input_shapes[0];
+            let output_shape = &output_shapes[0];
             assert!(kernel_shape.len() >= 2);
             assert!(kernel_shape[0] >= 0 && kernel_shape[1] >= 0);
 
@@ -554,11 +528,11 @@ pub fn compile(
             context.insert("alpha", &alpha);
             context.insert("beta", &beta);
 
-            if input_shape[0].dim(0) == 1 {
-                let threads = output_shape[0].dim(1);
+            if input_shapes[0].dim(0) == 1 {
+                let threads = output_shapes[0].dim(1);
                 ("matrix/gemm_1.wgsl".to_string(), threads as _, 1, 1)
             } else {
-                let threads = input_shape[0].dim(0) * input_shape[1].dim(1) / 16;
+                let threads = input_shapes[0].dim(0) * input_shapes[1].dim(1) / 16;
                 ("matrix/gemm.wgsl".to_string(), threads as _, 1, 1)
             }
         }
@@ -603,7 +577,7 @@ pub fn compile(
                     .iter()
                     .enumerate()
                     .map(|(i, x)| {
-                        let tmp = *x as f32 / input_shape[0].dim(i) as f32;
+                        let tmp = *x as f32 / input_shapes[0].dim(i) as f32;
                         format!("{:.2}", tmp)
                     })
                     .collect::<Vec<_>>()
@@ -663,12 +637,12 @@ pub fn compile(
         "Split" => {
             let mut axis = get_attribute("axis", Some(0), node)?;
             if axis < 0 {
-                axis += input_shape[0].element_count() as i64
+                axis += input_shapes[0].element_count() as i64
             }
             context.insert("axis", &axis);
 
-            let split_chunk = input_shape[0].dim(axis as usize) as usize / outputs.len();
-            let default_split = (1..=outputs.len())
+            let split_chunk = input_shapes[0].dim(axis as usize) as usize / output_shapes.len();
+            let default_split = (1..=output_shapes.len())
                 .map(|x| (x * split_chunk) as _)
                 .collect();
 
@@ -687,7 +661,7 @@ pub fn compile(
             let perms: Vec<i64> = get_attribute("perm", Some(default), node)?;
             let permuted_shapes = perms
                 .iter()
-                .map(|p| output_shape[0].dim(*p as usize))
+                .map(|p| output_shapes[0].dim(*p as usize))
                 .collect::<Vec<_>>();
 
             let mut chunks = vec![];
@@ -712,14 +686,78 @@ pub fn compile(
         .render(&template, &context)
         .expect("failed to render shader");
 
-    debug_assert!(
-        x < 16352,
-        "Node {} exceeds max compute by {}",
-        node.get_name(),
-        x - 16352
-    );
+    // Check if we remain within the limits of the thread count allowed by WebGPU
+    if x > MAX_COMPUTE_WORKGROUPS_PER_DIMENSION {
+        return Err(CompileError::ComputeLimitExceeded(
+            String::from("X threads"),
+            x as _,
+            MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
+        ));
+    }
+    if y > MAX_COMPUTE_WORKGROUPS_PER_DIMENSION {
+        return Err(CompileError::ComputeLimitExceeded(
+            String::from("Y threads"),
+            y as _,
+            MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
+        ));
+    }
+    if z > MAX_COMPUTE_WORKGROUPS_PER_DIMENSION {
+        return Err(CompileError::ComputeLimitExceeded(
+            String::from("Z threads"),
+            z as _,
+            MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
+        ));
+    }
+
     Ok(CompiledNode {
         shader,
         threads: (x, y, z),
     })
 }
+
+/// Determines the appropriate number of threads and workgroup size given a number of times the entry point of the shader should be run
+fn workgroup_size(
+    x: u64,
+    max_threads: u32,
+    max_workgroup_size: u32,
+) -> Result<(u32, u32), CompileError> {
+    let max_x = max_threads as u64;
+
+    Ok(if x > max_x {
+        let workgroup_size = ceil(x, max_x) as _;
+        let threads = ceil(x, workgroup_size as u64) as _;
+        log::info!(
+            "WGS: {} > {}, so workgroup size={} x threads={}",
+            x,
+            max_x,
+            workgroup_size,
+            threads
+        );
+
+        if threads > max_threads {
+            return Err(CompileError::ComputeLimitExceeded(
+                String::from("threads"),
+                threads,
+                max_threads,
+            ));
+        }
+
+        if workgroup_size > max_workgroup_size {
+            return Err(CompileError::ComputeLimitExceeded(
+                String::from("workgroup size"),
+                workgroup_size,
+                max_workgroup_size,
+            ));
+        }
+
+        log::info!(
+            "adjusting workgroup size to {}, threads to {} (was {})",
+            workgroup_size,
+            threads,
+            x
+        );
+        (threads, workgroup_size)
+    } else {
+        (x as u32, 1)
+    })
+}
diff --git a/src/gpu.rs b/src/gpu.rs
new file mode 100644
index 00000000..0e630093
--- /dev/null
+++ b/src/gpu.rs
@@ -0,0 +1,532 @@
+use std::{
+    borrow::Cow,
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
+
+use thiserror::Error;
+use wgpu::{Buffer, BufferUsages, CommandEncoder};
+
+use crate::{
+    compiler::{compile, CompileError, CompiledNode},
+    ir::{Node, NodeDefinition, NodeIdentifier, OperatorDefinition},
+    onnx::TensorProto,
+    resource::{self, resize},
+    utils::{ceil, Shape, MINIMUM_BUFFER_SIZE_BYTES},
+};
+
+/// The maximum number of bindings in a binding group (defined by wgpu)
+const MAX_BINDINGS_PER_GROUP: usize = 4;
+
+pub struct GpuModel {
+    device: wgpu::Device,
+    queue: wgpu::Queue,
+    onnx_opset_version: i64,
+    steps: Vec<GpuStep>,
+    inference_outputs: HashMap<String, InferenceOutput>,
+}
+
+/// An operation that is performed on the GPU as part of inference
+enum GpuStep {
+    /// A statically, pre-filled buffer containing tensor data
+    Initializer(Arc<Buffer>),
+
+    /// A buffer containing tensor data that is obtained from inference input
+    Input(String, Arc<Buffer>),
+
+    /// A GPU program (shader) that reads from buffers created by other steps and writes to output buffers
+    Operator {
+        pipeline: wgpu::ComputePipeline,
+        bind_groups: Vec<wgpu::BindGroup>,
+        threads: (u32, u32, u32),
+        output_tensors: Vec<GpuTensor>,
+    },
+
+    /// Operation that takes the output from a previous operation and assigns it to a second logical output
+    Forward(GpuTensor),
+
+    /// No-operation
+    None,
+}
+
+/// A tensor that resides in GPU memory
+#[derive(Clone)]
+struct GpuTensor {
+    buffer: Arc<Buffer>,
+    shape: Shape,
+}
+
+#[derive(Error, Debug)]
+pub enum GpuError {
+    #[error("compile error: {0}")]
+    CompileError(#[from] CompileError),
+
+    #[error("input not found: '{0}'")]
+    InputMissing(String),
+
+    #[error("node output not found: index {0}")]
+    OutputMissing(usize),
+}
+
+enum InferenceOutput {
+    InferenceInput(String),
+    Tensor(GpuTensor),
+}
+
+impl GpuModel {
+    /// Create a version of the specified model for which inference can be performed using the powers of the GPU
+    pub fn from(
+        root: Arc<Node>,
+        device: wgpu::Device,
+        queue: wgpu::Queue,
+        onnx_opset_version: i64,
+    ) -> Result<GpuModel, GpuError> {
+        let mut gpu_model = GpuModel {
+            device,
+            queue,
+            onnx_opset_version,
+            steps: vec![],
+            inference_outputs: HashMap::new(),
+        };
+
+        // Walk the IR DAG and encode into GPU execution steps
+        let mut readable_nodes: HashSet<NodeIdentifier> = HashSet::new();
+        let mut node_outputs = HashMap::<NodeIdentifier, Vec<GpuTensor>>::new();
+        gpu_model.sequence(root.clone(), &mut readable_nodes, &mut node_outputs)?;
+
+        // Find out which outputs we should return as inference outputs
+        if let NodeDefinition::Outputs { names } = &root.definition {
+            for (usize, output_name) in names.iter().enumerate() {
+                let input = &root.inputs[usize];
+                gpu_model.inference_outputs.insert(
+                    output_name.to_string(),
+                    match &input.source_node.definition {
+                        NodeDefinition::Operator(_) | NodeDefinition::Tensor(_) => {
+                            let source_identifier = input.source_node.identifier();
+                            let outputs = &node_outputs[&source_identifier];
+                            let tensor = outputs[input.output_index].clone();
+                            InferenceOutput::Tensor(tensor)
+                        }
+                        NodeDefinition::Input(proto) => {
+                            InferenceOutput::InferenceInput(proto.get_name().to_string())
+                        }
+                        NodeDefinition::Outputs { .. } => {
+                            unimplemented!("output after output node")
+                        }
+                        NodeDefinition::Missing => {
+                            unimplemented!("optional input after output node")
+                        }
+                    },
+                );
+            }
+        } else {
+            unimplemented!("reading from non-outputs IR node")
+        }
+
+        // Upload the data (for initializers etc.) by submitting an empty command queue
+        log::info!("submit initializer buffers");
+        let encoder = gpu_model
+            .device
+            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
+        gpu_model.queue.submit(Some(encoder.finish()));
+
+        Ok(gpu_model)
+    }
+
+    /// Write commands to the GPU to create the necessary resources to be able to perform inference (e.g. allocates buffers
+    /// for intermediate results, compiles shader code, determines which outputs to return, etc.).
+    fn sequence<'model>(
+        &mut self,
+        node: Arc<Node<'model>>,
+        nodes_readable: &mut HashSet<NodeIdentifier<'model>>,
+        node_outputs: &mut HashMap<NodeIdentifier<'model>, Vec<GpuTensor>>,
+    ) -> Result<(), GpuError> {
+        let node_identifier = node.identifier();
+        let outputs_readable = nodes_readable.contains(&node_identifier);
+
+        // Sequence inputs
+        let mut input_tensors: Vec<GpuTensor> = vec![];
+        for node_input in &node.inputs {
+            // If this node is an output node, mark input nodes as 'readable', meaning that their output buffers need to be created as readable buffers
+            if let NodeDefinition::Outputs { .. } = &node.definition {
+                nodes_readable.insert(node_input.source_node.identifier());
+            }
+
+            if let NodeDefinition::Operator(op_def) = &node.definition {
+                // For these ops we just forward the buffer (so we should also forward readability)
+                if op_def.proto.get_op_type() == "Reshape" {
+                    nodes_readable.insert(node_input.source_node.identifier());
+                }
+            }
+
+            // Sequence the source node
+            self.sequence(node_input.source_node.clone(), nodes_readable, node_outputs)?;
+
+            // Select the tensor we want for our input
+            let source_identifier = node_input.source_node.identifier();
+            let input_tensor = {
+                let source_outputs = &node_outputs[&source_identifier];
+                source_outputs
+                    .get(node_input.output_index)
+                    .ok_or(GpuError::OutputMissing(node_input.output_index))?
+                    .clone()
+            };
+            input_tensors.push(input_tensor);
+        }
+
+        // Sequence self
+        if let std::collections::hash_map::Entry::Vacant(e) = node_outputs.entry(node_identifier) {
+            log::info!(
+                "sequence {:?} (outputs readable={:?})",
+                node.definition,
+                outputs_readable
+            );
+
+            let mut output_tensors = vec![];
+            let gpu_op: GpuStep = match &node.definition {
+                NodeDefinition::Operator(op_def) => {
+                    let gpu_op = op_def.gpu_op(
+                        &self.device,
+                        outputs_readable,
+                        self.onnx_opset_version,
+                        &input_tensors,
+                    )?;
+
+                    match &gpu_op {
+                        GpuStep::Operator {
+                            output_tensors: op_output_tensors,
+                            ..
+                        } => {
+                            output_tensors.extend(op_output_tensors.iter().cloned());
+                        }
+                        GpuStep::Forward(output_tensor) => {
+                            output_tensors.push(output_tensor.clone());
+                        }
+                        _ => unreachable!("gpu_op for operator produced something unexpected"),
+                    }
+
+                    gpu_op
+                }
+                NodeDefinition::Tensor(tensor_def) => {
+                    let tensor_buffer = Arc::new(tensor_def.buffer(&self.device, outputs_readable));
+                    output_tensors.push(GpuTensor {
+                        shape: Shape::from(tensor_def.get_dims()),
+                        buffer: tensor_buffer.clone(),
+                    });
+                    GpuStep::Initializer(tensor_buffer)
+                }
+                NodeDefinition::Input(input_def) => {
+                    if outputs_readable {
+                        log::warn!(
+                            "it looks like you will be reading back an inference input as output"
+                        );
+                    }
+
+                    let input_shape = input_def.get_shape();
+                    let input_buffer = Arc::new(resource::buffer(
+                        &self.device,
+                        input_shape.buffer_len() as _,
+                        input_def.get_name(),
+                        BufferUsages::STORAGE | BufferUsages::COPY_DST,
+                    ));
+
+                    output_tensors.push(GpuTensor {
+                        shape: input_shape,
+                        buffer: input_buffer.clone(),
+                    });
+
+                    GpuStep::Input(input_def.get_name().to_string(), input_buffer)
+                }
+                NodeDefinition::Missing | NodeDefinition::Outputs { .. } => {
+                    // Nothing to sequence
+                    GpuStep::None
+                }
+            };
+
+            e.insert(output_tensors);
+            self.steps.push(gpu_op);
+            Ok(())
+        } else {
+            // This node is already sequenced
+            log::debug!("not sequencing (already seen) {:?}", node.definition);
+            Ok(())
+        }
+    }
+
+    /// Perform inference using this model and the specified inference inputs.
+    pub async fn infer(
+        &self,
+        inference_inputs: &HashMap<String, &[f32]>,
+    ) -> Result<HashMap<String, Vec<f32>>, GpuError> {
+        log::info!("encode inference steps");
+        let mut encoder = self
+            .device
+            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
+        for step in &self.steps {
+            step.encode(&self.queue, &mut encoder, inference_inputs)?;
+        }
+        log::info!("submit inference steps");
+        self.queue.submit(Some(encoder.finish()));
+        log::info!("inference completed");
+        self.read_outputs(inference_inputs).await
+    }
+
+    /// Reads the relevant buffers for the requested inference outputs
+    async fn read_outputs(
+        &self,
+        inference_inputs: &HashMap<String, &[f32]>,
+    ) -> Result<HashMap<String, Vec<f32>>, GpuError> {
+        let mut output_data = HashMap::new();
+
+        for (output_name, output_source) in &self.inference_outputs {
+            output_data.insert(
+                output_name.to_string(),
+                match output_source {
+                    InferenceOutput::InferenceInput(input_name) => {
+                        inference_inputs[input_name].to_vec()
+                    }
+                    InferenceOutput::Tensor(tensor) => tensor.read_to_vec(&self.device).await?,
+                },
+            );
+        }
+
+        Ok(output_data)
+    }
+}
+
+trait TensorProtoExtra {
+    fn buffer(&self, device: &wgpu::Device, readable: bool) -> Buffer;
+}
+
+impl TensorProtoExtra for TensorProto {
+    /// Create a GPU buffer containing the data of this initializer
+    fn buffer(&self, device: &wgpu::Device, readable: bool) -> Buffer {
+        let input_shape = Shape::from(self.get_dims());
+        log::info!(
+            "creating tensor buffer {} shape {}",
+            self.get_name(),
+            input_shape
+        );
+
+        let data = self.get_float_data();
+        let raw_data = if !data.is_empty() {
+            bytemuck::cast_slice(data)
+        } else {
+            self.get_raw_data()
+        };
+
+        let buffer_usage = match readable {
+            true => BufferUsages::MAP_READ | BufferUsages::STORAGE,
+            false => BufferUsages::STORAGE,
+        };
+
+        // Do not create buffers that are too small
+        if raw_data.len() < MINIMUM_BUFFER_SIZE_BYTES as _ {
+            let mut larger_raw_data = raw_data.to_vec();
+            larger_raw_data.resize(MINIMUM_BUFFER_SIZE_BYTES as _, 0);
+            resource::create_buffer_init(device, &larger_raw_data, self.get_name(), buffer_usage)
+        } else {
+            resource::create_buffer_init(device, raw_data, self.get_name(), buffer_usage)
+        }
+    }
+}
+
+impl<'model> OperatorDefinition<'model> {
+    fn gpu_op(
+        &self,
+        device: &wgpu::Device,
+        outputs_readable: bool,
+        opset_version: i64,
+        input_tensors: &[GpuTensor],
+    ) -> Result<GpuStep, GpuError> {
+        let proto = &self.proto;
+
+        // Some nodes have specific GPU implementations, match these here
+        match proto.get_op_type() {
+            // Some ops do nothing but forward their input
+            "Reshape" | "Identity" | "Flatten" | "Squeeze" | "Unsqueeze" | "Dropout" => {
+                let value_shape = &self.output_shapes[0];
+                let output_tensor = GpuTensor {
+                    buffer: input_tensors[0].buffer.clone(),
+                    shape: value_shape.clone(),
+                };
+                return Ok(GpuStep::Forward(output_tensor));
+            }
+            _ => {}
+        }
+
+        let label = Some(proto.get_name());
+
+        // Create output buffers for this op node
+        let output_tensors: Vec<GpuTensor> = proto
+            .get_output()
+            .iter()
+            .enumerate()
+            .map(|(output_index, output_name)| {
+                let value_shape = &self.output_shapes[output_index];
+                log::info!(
+                    "Creating op output buffer for output #{} ({}) of {} shaped {}",
+                    output_index,
+                    output_name,
+                    proto.get_name(),
+                    value_shape
+                );
+
+                let buffer_usage = if outputs_readable {
+                    BufferUsages::STORAGE | BufferUsages::MAP_READ
+                } else {
+                    BufferUsages::STORAGE
+                };
+
+                let buffer = Arc::new(resource::buffer(
+                    device,
+                    value_shape.buffer_len() as _,
+                    output_name.as_str(),
+                    buffer_usage,
+                ));
+                GpuTensor {
+                    buffer,
+                    shape: value_shape.clone(),
+                }
+            })
+            .collect();
+
+        let input_shapes: Vec<&Shape> = input_tensors.iter().map(|input| &input.shape).collect();
+        let output_shapes: Vec<&Shape> = self.output_shapes.iter().collect();
+
+        // Compile shader for node
+        let CompiledNode { shader, threads } =
+            compile(proto, &input_shapes, &output_shapes, opset_version)?;
+        log::debug!("shader: {}", shader);
+
+        // Bind input and output buffers to the shader
+        let mut binding_counter: usize = 0;
+        let mut entries = vec![];
+
+        // Bind input buffers
+        for input in input_tensors {
+            // Bindings are numbered 0...3 (MAX_BINDINGS_PER_GROUP-1) in binding groups (starting at group 0)
+            let binding_index = (binding_counter % MAX_BINDINGS_PER_GROUP) as u32;
+
+            entries.push(wgpu::BindGroupEntry {
+                binding: binding_index,
+                resource: input.buffer.as_entire_binding(),
+            });
+            binding_counter += 1;
+        }
+
+        // Bind output buffers
+        for output_tensor in &output_tensors {
+            // Bindings are numbered 0...3 (MAX_BINDINGS_PER_GROUP-1) in binding groups (starting at group 0)
+            let binding_index = (binding_counter % MAX_BINDINGS_PER_GROUP) as u32;
+
+            entries.push(wgpu::BindGroupEntry {
+                binding: binding_index,
+                resource: output_tensor.buffer.as_entire_binding(),
+            });
+            binding_counter += 1;
+        }
+
+        // Set up a pipeline (basically the shader source code with some metadata that determines how it will be executed)
+        let mut bind_groups = vec![];
+        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+            label,
+            layout: None,
+            module: &device.create_shader_module(&wgpu::ShaderModuleDescriptor {
+                label,
+                source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(&shader)),
+            }),
+            entry_point: "main",
+        });
+
+        // Create 'bind groups' (groups of bound buffers)
+        let number_of_groups = ceil(binding_counter as u64, MAX_BINDINGS_PER_GROUP as u64) as usize;
+        for group_index in 0..number_of_groups {
+            let group_range = group_index * MAX_BINDINGS_PER_GROUP
+                ..usize::min(
+                    binding_counter as _,
+                    (group_index + 1) * MAX_BINDINGS_PER_GROUP,
+                );
+            bind_groups.push(device.create_bind_group(&wgpu::BindGroupDescriptor {
+                label,
+                layout: &pipeline.get_bind_group_layout(group_index as u32),
+                entries: &entries[group_range],
+            }));
+        }
+
+        Ok(GpuStep::Operator {
+            output_tensors,
+            pipeline,
+            bind_groups,
+            threads,
+        })
+    }
+}
+
+impl GpuStep {
+    /// Writes the necessary commands for the GPU to execute this step into the command queue. Among other things this means
+    /// writing the inference input data to the appropriate (already created) buffers.
+    fn encode(
+        &self,
+        queue: &wgpu::Queue,
+        encoder: &mut CommandEncoder,
+        inputs: &HashMap<String, &[f32]>,
+    ) -> Result<(), GpuError> {
+        match self {
+            GpuStep::None | GpuStep::Forward(_) | GpuStep::Initializer(_) => {
+                // Buffer already filled, no need to encode anything at this point.
+                Ok(())
+            }
+            GpuStep::Input(input_name, input_buffer) => {
+                // Encode a command to write the input data to the corresponding input buffer (which was created empty
+                // by `GpuModel::from`
+                let input_data = inputs
+                    .get(input_name)
+                    .ok_or_else(|| GpuError::InputMissing(input_name.to_string()))?;
+                log::info!("- write input data for {}", input_name);
+                queue.write_buffer(
+                    input_buffer,
+                    0,
+                    bytemuck::cast_slice(&resize(input_data.to_vec())),
+                );
+                Ok(())
+            }
+            GpuStep::Operator {
+                pipeline,
+                bind_groups,
+                threads,
+                ..
+            } => {
+                // Encode a command for invocation of a shader.
+                let mut compute_pass =
+                    encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None });
+                compute_pass.set_pipeline(pipeline);
+                for (index, bind_group) in bind_groups.iter().enumerate() {
+                    compute_pass.set_bind_group(index as u32, bind_group, &[]);
+                }
+                let (x, y, z) = *threads;
+                compute_pass.dispatch(x, y, z);
+                Ok(())
+            }
+        }
+    }
+}
+
+impl GpuTensor {
+    /// Read the tensor from GPU memory to main memory (as Vec<f32>)
+    async fn read_to_vec(&self, device: &wgpu::Device) -> Result<Vec<f32>, GpuError> {
+        let buffer_slice = self.buffer.slice(..);
+        let buffer_future = buffer_slice.map_async(wgpu::MapMode::Read);
+        device.poll(wgpu::Maintain::Wait);
+        buffer_future.await.expect("failed to run compute on gpu!");
+        let output_data = buffer_slice.get_mapped_range();
+
+        // The actual buffer may be bigger than what we should return, because buffers have a minimum size in wgpu
+        // Fetch the size we should expect so we can chop the buffer to the correct size
+        let output_buffer_size = self.shape.element_count() as usize;
+        let result = bytemuck::cast_slice(&output_data)[..output_buffer_size].to_vec();
+        drop(output_data);
+        self.buffer.unmap();
+        Ok(result)
+    }
+}
diff --git a/src/ir.rs b/src/ir.rs
new file mode 100644
index 00000000..b3d56ab2
--- /dev/null
+++ b/src/ir.rs
@@ -0,0 +1,334 @@
+use crate::onnx::{ModelProto, NodeProto, TensorProto, ValueInfoProto};
+use crate::utils::Shape;
+use std::borrow::Cow;
+use std::fmt::Debug;
+use std::hash::Hash;
+use std::ptr;
+use std::{collections::HashMap, sync::Arc};
+use thiserror::Error;
+
+#[derive(Clone)]
+pub struct OperatorDefinition<'model> {
+    pub(crate) proto: Cow<'model, NodeProto>,
+    pub(crate) output_shapes: Vec<Shape>,
+}
+
+impl<'model> OperatorDefinition<'model> {
+    pub fn from(
+        node: Cow<'model, NodeProto>,
+        value_shapes: &HashMap<&'model str, Shape>,
+    ) -> Result<OperatorDefinition<'model>, IrError> {
+        let mut output_shapes: Vec<Shape> = Vec::with_capacity(node.get_output().len());
+        for output_name in node.get_output() {
+            if !value_shapes.contains_key(output_name.as_str()) {
+                return Err(IrError::OutputNodeNotFound(output_name.to_string()));
+            }
+
+            output_shapes.push(value_shapes[&output_name.as_str()].clone());
+        }
+        Ok(OperatorDefinition {
+            proto: node,
+            output_shapes,
+        })
+    }
+}
+
+#[derive(Clone)]
+pub enum NodeDefinition<'model> {
+    Operator(Box<OperatorDefinition<'model>>),
+    Tensor(Box<Cow<'model, TensorProto>>),
+    Input(&'model ValueInfoProto),
+    Outputs { names: Vec<&'model str> },
+    Missing, // A missing input (optional)
+}
+
+static MISSING_OPTIONAL_INPUT: NodeDefinition<'static> = NodeDefinition::Missing;
+
+#[derive(Clone)]
+pub struct Input<'model> {
+    pub source_node: Arc<Node<'model>>,
+    pub output_index: usize,
+}
+
+pub struct Node<'model> {
+    pub definition: NodeDefinition<'model>,
+    pub inputs: Vec<Input<'model>>,
+}
+
+#[derive(Debug, Error)]
+pub enum IrError {
+    #[error("output node for output {0} not found")]
+    OutputNodeNotFound(String),
+
+    #[error("could not find node corresponding to input {input_name} of node {target_node_name}")]
+    InputNodeNotFound {
+        target_node_name: String,
+        input_name: String,
+    },
+}
+
+impl<'m> NodeDefinition<'m> {
+    pub fn get_name(&self) -> Cow<'_, str> {
+        match self {
+            NodeDefinition::Operator(op_def) => Cow::from(op_def.proto.get_name()),
+            NodeDefinition::Tensor(t) => Cow::from(t.get_name()),
+            NodeDefinition::Input(i) => Cow::from(i.get_name()),
+            NodeDefinition::Outputs { .. } => Cow::from(" "),
+            NodeDefinition::Missing => Cow::from(""),
+        }
+    }
+
+    pub fn output_name(&self, output_index: usize) -> Cow<'_, str> {
+        match self {
+            NodeDefinition::Operator(op_def) => {
+                Cow::Borrowed(&op_def.proto.get_output()[output_index])
+            }
+            NodeDefinition::Tensor(proto) => Cow::from(proto.get_name()),
+            NodeDefinition::Input(proto) => Cow::from(proto.get_name()),
+            NodeDefinition::Outputs { .. } => panic!("can't get output name for outputs node"),
+            NodeDefinition::Missing => panic!("can't get output name for missing node"),
+        }
+    }
+}
+
+impl<'model> Node<'model> {
+    pub fn new(variant: NodeDefinition<'model>) -> Node<'model> {
+        Node {
+            definition: variant,
+            inputs: vec![],
+        }
+    }
+
+    pub fn definition(&self) -> &NodeDefinition<'model> {
+        &self.definition
+    }
+
+    /// Construct part of the intermediate representation tree for the indicated node.
+    pub fn from_node<'a>(
+        model: &'model ModelProto,
+        node: Cow<'model, NodeProto>,
+        value_shapes: &HashMap<&'model str, Shape>,
+        node_definitions_by_output: &'a HashMap<String, NodeDefinition<'model>>,
+        nodes_by_name: &mut HashMap<String, Arc<Node<'model>>>,
+    ) -> Result<Arc<Node<'model>>, IrError> {
+        let node_name = node.get_name();
+        // Did we already translate this node before?
+        if nodes_by_name.contains_key(node_name) {
+            let n = nodes_by_name.get(node_name).unwrap();
+            return Ok(n.clone());
+        }
+
+        let inputs: Result<Vec<Input<'model>>, IrError> = node
+            .get_input()
+            .iter()
+            .map(|input_name: &'model String| {
+                let source_node_definition = node_definitions_by_output
+                    .get(&input_name.to_string())
+                    .unwrap_or(&MISSING_OPTIONAL_INPUT);
+
+                Ok(match source_node_definition {
+                    // The source is another op - continue translating that node
+                    NodeDefinition::Operator(source_node_proto) => Input {
+                        source_node: Node::from_node(
+                            model,
+                            source_node_proto.proto.clone(),
+                            value_shapes,
+                            node_definitions_by_output,
+                            nodes_by_name,
+                        )?,
+                        output_index: source_node_proto
+                            .proto
+                            .get_output()
+                            .iter()
+                            .position(|s| s == input_name)
+                            .ok_or_else(|| IrError::OutputNodeNotFound(input_name.to_string()))?,
+                    },
+                    _ => {
+                        // The source is an initializer or model onput
+                        let source_name = source_node_definition.get_name();
+
+                        Input {
+                            output_index: 0,
+                            // Did we already translate this node?
+                            source_node: match nodes_by_name.get(&source_name.to_string()) {
+                                Some(node) => node.clone(),
+                                None => {
+                                    let node = Arc::new(Node::new(source_node_definition.clone()));
+                                    nodes_by_name.insert(source_name.into(), node.clone());
+                                    node
+                                }
+                            },
+                        }
+                    }
+                })
+            })
+            .collect();
+
+        let translated = Arc::new(Node {
+            definition: NodeDefinition::Operator(Box::new(OperatorDefinition::from(
+                node.clone(),
+                value_shapes,
+            )?)),
+            inputs: inputs?,
+        });
+        nodes_by_name.insert(node_name.to_string(), translated.clone());
+        Ok(translated)
+    }
+
+    /// Construct an intermediate representation graph for calculating the output with the specified name.
+    pub fn from_model(model: &'model ModelProto) -> Result<Arc<Node<'model>>, IrError> {
+        // Collect value shapes
+        let mut value_shapes: HashMap<&'model str, Shape> = HashMap::new();
+        for vi in model.get_graph().get_value_info() {
+            value_shapes.insert(vi.get_name(), vi.get_shape());
+        }
+
+        for vi in model.get_graph().get_output() {
+            let output_name = vi.get_name();
+            if !output_name.is_empty() {
+                value_shapes.insert(output_name, vi.get_shape());
+            }
+        }
+
+        // Sort nodes by output nodes
+        let mut node_definitions_by_output = HashMap::<String, NodeDefinition<'model>>::new();
+        for node in model.get_graph().get_node().iter() {
+            let node_def = NodeDefinition::Operator(Box::new(OperatorDefinition::from(
+                Cow::Borrowed(node),
+                &value_shapes,
+            )?));
+            for output in node.get_output() {
+                if !output.is_empty() {
+                    node_definitions_by_output.insert(output.to_string(), node_def.clone());
+                }
+            }
+        }
+
+        // Collect intializer info
+        for initializer in model.get_graph().get_initializer().iter() {
+            log::info!("Initializer {}", initializer.get_name());
+            node_definitions_by_output.insert(
+                initializer.get_name().to_string(),
+                NodeDefinition::Tensor(Box::new(Cow::Borrowed(initializer))),
+            );
+        }
+
+        // Collect input name
+        for input in model.get_graph().get_input().iter() {
+            if !node_definitions_by_output.contains_key(input.get_name()) {
+                log::info!("Input {}", input.get_name());
+                node_definitions_by_output
+                    .insert(input.get_name().to_string(), NodeDefinition::Input(input));
+            } else {
+                log::info!(
+                    "Skipping input definition {}: already defined",
+                    input.get_name()
+                );
+            }
+        }
+
+        let mut nodes_by_name = HashMap::new();
+
+        let output_nodes: Result<Vec<Input<'model>>, IrError> = model
+            .get_graph()
+            .get_output()
+            .iter()
+            .map(|output_def| {
+                let output_name_string = output_def.get_name().to_string();
+                let output_node = model
+                    .get_graph()
+                    .get_node()
+                    .iter()
+                    .find(|x| -> bool { x.get_output().contains(&output_name_string) })
+                    .ok_or(IrError::OutputNodeNotFound(output_name_string))?;
+
+                let source_node = Node::<'model>::from_node(
+                    model,
+                    Cow::Borrowed(output_node),
+                    &value_shapes,
+                    &node_definitions_by_output,
+                    &mut nodes_by_name,
+                )?;
+
+                let output_index = output_node
+                    .get_output()
+                    .iter()
+                    .position(|s| s == output_def.get_name())
+                    .ok_or_else(|| {
+                        IrError::OutputNodeNotFound(output_def.get_name().to_string())
+                    })?;
+
+                Ok(Input {
+                    source_node,
+                    output_index,
+                })
+            })
+            .collect();
+
+        let output_names: Vec<&str> = model
+            .get_graph()
+            .get_output()
+            .iter()
+            .map(|output_def| output_def.get_name())
+            .collect();
+
+        Ok(Arc::new(Node {
+            definition: NodeDefinition::Outputs {
+                names: output_names,
+            },
+            inputs: output_nodes?,
+        }))
+    }
+
+    pub fn output_shape(&self, output_index: usize) -> Shape {
+        match (&self.definition, output_index) {
+            (NodeDefinition::Operator(op_def), index) => op_def.output_shapes[index].clone(),
+            (NodeDefinition::Tensor(tensor_proto), 0) => Shape::from(tensor_proto.get_dims()),
+            (NodeDefinition::Input(input_proto), 0) => input_proto.get_shape(),
+            (NodeDefinition::Outputs { .. }, _) => panic!("output node has no outputs!"),
+            (_, _) => panic!("node has no output at index {}", output_index),
+        }
+    }
+}
+
+impl<'model> Debug for NodeDefinition<'model> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            NodeDefinition::Operator(def) => {
+                write!(
+                    f,
+                    "op: {} ({})",
+                    def.proto.get_name(),
+                    def.proto.get_op_type()
+                )
+            }
+            NodeDefinition::Tensor(def) => write!(f, "tensor {}", def.get_name()),
+            NodeDefinition::Input(def) => write!(f, "input {}", def.get_name()),
+            NodeDefinition::Outputs { .. } => write!(f, "outputs"),
+            NodeDefinition::Missing => write!(f, "missing (optional)"),
+        }
+    }
+}
+
+/// Wrap an Arc<Node> in a struct so we can implement pointer-based comparison for it, and use them as keys in a HashSet/HashMap
+pub struct NodeIdentifier<'model>(Arc<Node<'model>>);
+
+impl<'model> Hash for NodeIdentifier<'model> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        ptr::hash(Arc::as_ptr(&self.0), state)
+    }
+}
+
+impl<'model> PartialEq for NodeIdentifier<'model> {
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.0, &other.0)
+    }
+}
+
+impl<'model> Eq for NodeIdentifier<'model> {}
+
+impl<'model> Node<'model> {
+    pub fn identifier(self: &Arc<Self>) -> NodeIdentifier<'model> {
+        NodeIdentifier(self.clone())
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index d7e80d11..6b05e02e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,21 +1,23 @@
 pub mod compiler;
+mod gpu;
+pub mod ir;
 pub mod onnx;
-pub mod optimisation;
+pub mod optimizer;
 pub mod resource;
-pub mod sequencer;
 pub mod utils;
 
 #[macro_use]
 extern crate lazy_static;
 
 use compiler::CompileError;
-use onnx::ValueInfoProto;
-use optimisation::{OptimizationError, OptimizedModel};
+use gpu::GpuError;
+use ir::IrError;
+use optimizer::{Optimizer, OptimizerError};
 use protobuf::{self, Message, ProtobufError};
 use std::collections::HashMap;
 use std::result::Result;
 
-use crate::resource::resize;
+use crate::gpu::GpuModel;
 use thiserror::Error;
 
 #[derive(Error, Debug)]
@@ -38,32 +40,8 @@ pub enum WonnxError {
 /// ```ignore
 /// let mut session = Session::from_path("path/to/model.onnx").await.unwrap();
 /// ```
-// +----------------+
-// |ONNX Path       |
-// +----------------+
-//         v
-// +----------------+
-// |ONNX Model      |
-// +----------------+
-//         v
-// +----------------+
-// |Session         |                        Optimisation on       Node -> WGSL Shader
-// +----------------+                        Mutiple node
-//         v
-// +----------------+   +----------------+   +----------------+   +----------------+
-// |load (once)     | > |Load params     | > |Sequencer       | > |Compiler        |
-// +----------------+   +----------------+   +----------------+   +----------------+
-//         v
-// +----------------+
-// |run             |  (Can be run multiple times)
-// +----------------+
-//
-//
 pub struct Session {
-    pub device: wgpu::Device,
-    pub queue: wgpu::Queue,
-    pub optimized_model: OptimizedModel,
-    pub outputs: Vec<ValueInfoProto>,
+    gpu_model: GpuModel,
 }
 
 #[derive(Error, Debug)]
@@ -84,30 +62,34 @@ pub enum SessionError {
     )]
     InvalidOutput(String),
 
-    #[error("optimization failed: {0}")]
-    OptimizationFailed(#[from] OptimizationError),
-
     #[error("more than one ONNX opset was specified: {0} and {1}")]
-    DuplicateONNXOpset(i64, i64),
+    DuplicateOnnxOpset(i64, i64),
 
     #[error("the model references an unknown opset: '{0}'")]
     UnknownOpset(String),
 
     #[error("the model did not reference a specific version of the ONNX opset")]
-    UnknownONNXOpsetVersion,
+    UnknownOnnxOpsetVersion,
+
+    #[error("IR error: {0}")]
+    IrError(#[from] IrError),
+
+    #[error("GPU model error: {0}")]
+    GpuError(#[from] GpuError),
+
+    #[error("optimizer error: {0}")]
+    OptimizerError(#[from] OptimizerError),
 }
 
 impl Session {
     // Read an ONNX model from a path and create a session.
     pub async fn from_path(path: &str) -> Result<Session, SessionError> {
         let model = onnx::ModelProto::parse_from_bytes(&std::fs::read(path)?)?;
-
         Session::from_model(model).await
     }
 
     pub async fn from_bytes(bytes: &[u8]) -> Result<Session, SessionError> {
         let model = onnx::ModelProto::parse_from_bytes(bytes)?;
-
         Session::from_model(model).await
     }
 
@@ -125,7 +107,7 @@ impl Session {
                     // This is a reference to the ONNX specification op set
                     if let Some(onnx_version) = onnx_opset_version {
                         if opset_import.get_version() != onnx_version {
-                            return Err(SessionError::DuplicateONNXOpset(
+                            return Err(SessionError::DuplicateOnnxOpset(
                                 onnx_version,
                                 opset_import.get_version(),
                             ));
@@ -141,97 +123,20 @@ impl Session {
         }
 
         // Optimize and compile the model graph to a set of buffers and 'builders' which can basically run GPU shader code referencing these buffers
-        let onnx_opset_version = onnx_opset_version.ok_or(SessionError::UnknownONNXOpsetVersion)?;
-        let graph = model.get_graph();
-        let optimized_model = optimisation::load(graph, &device, onnx_opset_version)?;
-
-        // The data is loaded after the first submit
-        let encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
-        queue.submit(Some(encoder.finish()));
-
-        Ok(Session {
-            device,
-            queue,
-            optimized_model,
-            outputs: graph.get_output().to_vec(),
-        })
+        let onnx_opset_version = onnx_opset_version.ok_or(SessionError::UnknownOnnxOpsetVersion)?;
+
+        let mut optimizer = Optimizer::new();
+        let ir = optimizer.optimize(ir::Node::from_model(&model)?)?;
+        let gpu_model = GpuModel::from(ir, device, queue, onnx_opset_version)?;
+
+        Ok(Session { gpu_model })
     }
 
-    // Run use the element loaded into the session to produce the inference.
-    // It copy input data to the buffers.
-    // Run the command encoder.
-    // Copy the output into an exit buffer that can be deleted.
+    /// Perform inference given the inputs provided and return all the outputs the model was compiled to return.
     pub async fn run(
         &self,
-        inputs: HashMap<String, &[f32]>,
+        inputs: &HashMap<String, &[f32]>,
     ) -> Result<HashMap<String, Vec<f32>>, SessionError> {
-        let device = &self.device;
-        let queue = &self.queue;
-        let buffers = &self.optimized_model.buffers;
-        let outputs = &self.outputs;
-
-        // Copy input data to the GPU buffer where our first node will read it
-        for (input, data) in inputs {
-            queue.write_buffer(
-                buffers
-                    .get(&input)
-                    .ok_or(SessionError::InvalidInput(input))?,
-                0,
-                bytemuck::cast_slice(&resize(data.to_vec())),
-            )
-        }
-
-        // Run the command encoder generated during the load
-        let mut encoder =
-            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
-
-        // Sequentially execute each node (compiled to a shader 'builder')
-        for builder in &self.optimized_model.builders {
-            let mut cpass =
-                encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None });
-            cpass.set_pipeline(&builder.pipeline);
-            for (index, bind_group) in builder.bind_groups.iter().enumerate() {
-                cpass.set_bind_group(index as u32, bind_group, &[]);
-            }
-
-            let (x, y, z) = builder.threads;
-            cpass.dispatch(x, y, z); // Number of cells to run, the (x,y,z) size of item being processed
-        }
-
-        queue.submit(Some(encoder.finish()));
-        let mut results = HashMap::new();
-
-        for output in outputs {
-            let output_name = output.get_name();
-
-            // Copy the output data into the staging buffer.
-            let buffer = buffers
-                .get(output_name)
-                .ok_or_else(|| SessionError::InvalidOutput(output_name.to_string()))?;
-
-            let buffer_slice = buffer.slice(..);
-            // TODO: Define behavior for multi output.
-            let buffer_future = buffer_slice.map_async(wgpu::MapMode::Read);
-
-            device.poll(wgpu::Maintain::Wait);
-
-            buffer_future.await.expect("failed to run compute on gpu!");
-            // Gets contents of buffer
-            let data = buffer_slice.get_mapped_range();
-            // Since contents are got in bytes, this converts these bytes back to f32
-
-            // The actual buffer may be bigger than what we should return, because buffers have a minimum size in wgpu
-            // Fetch the size we should expect so we can chop the buffer to the correct size
-            let output_buffer_size = output.get_shape().element_count() as usize;
-
-            results.insert(
-                output_name.to_string(),
-                bytemuck::cast_slice(&data)[..output_buffer_size].to_vec(),
-            );
-            drop(data);
-            buffer.unmap();
-        }
-
-        Ok(results)
+        Ok(self.gpu_model.infer(inputs).await?)
     }
 }
diff --git a/src/optimisation.rs b/src/optimisation.rs
deleted file mode 100644
index e5e0e5d0..00000000
--- a/src/optimisation.rs
+++ /dev/null
@@ -1,190 +0,0 @@
-use crate::{
-    compiler::{compile, CompileError, CompiledNode},
-    resource,
-    sequencer::{sequence, SequenceError},
-    utils::{ceil, dimensions_infos, initializers},
-    Result,
-};
-
-use std::{borrow::Cow, collections::HashMap};
-
-use log::info;
-use thiserror::Error;
-use wgpu::BufferUsages;
-
-pub struct EncoderBuilder {
-    pub pipeline: wgpu::ComputePipeline,
-    pub bind_groups: Vec<wgpu::BindGroup>,
-    pub threads: (u32, u32, u32),
-}
-
-const MAX_BINDINGS_PER_GROUP: usize = 4;
-
-#[derive(Error, Debug)]
-pub enum OptimizationError {
-    #[error("compilation failed: {0}")]
-    CompilationFailed(#[from] CompileError),
-
-    #[error("output shape was not provided. You can use python's onnx-simplifier to generate implied shapes.")]
-    OutputShapeMissing,
-
-    #[error("tensor metadata missing: '{0}'")]
-    TensorMetadataMissing(String),
-
-    #[error("sequencing failed: {0}")]
-    SequencingFailed(#[from] SequenceError),
-}
-
-pub struct OptimizedModel {
-    pub buffers: HashMap<String, wgpu::Buffer>,
-    pub builders: Vec<EncoderBuilder>,
-}
-
-pub fn load(
-    graph: &crate::onnx::GraphProto,
-    device: &wgpu::Device,
-    opset_version: i64,
-) -> Result<OptimizedModel, OptimizationError> {
-    let initializers = initializers(graph);
-    let mut shapes_info = dimensions_infos(graph);
-    let mut buffers = HashMap::new();
-
-    for (input_name, input_shape) in shapes_info.iter() {
-        buffers.insert(
-            input_name.clone(),
-            resource::buffer(
-                device,
-                input_shape.buffer_len() as _,
-                input_name,
-                BufferUsages::STORAGE | BufferUsages::COPY_DST,
-            ),
-        );
-    }
-
-    let base_nodes = graph.get_node();
-    let n = base_nodes.len();
-    let mut node_index = 0;
-    let mut builders = vec![];
-
-    let output_info = &graph.get_output().to_vec();
-
-    while node_index < n {
-        let nodes = &base_nodes[node_index..];
-        let names = nodes
-            .iter()
-            .map(|node| node.get_op_type())
-            .collect::<Vec<_>>();
-
-        // Generate the shader source code for this node
-        let sequence = sequence(
-            &names,
-            nodes,
-            device,
-            &initializers,
-            &mut buffers,
-            &mut shapes_info,
-        )?;
-
-        let current_node = sequence.node;
-        let CompiledNode { shader, threads } = compile(&current_node, &shapes_info, opset_version)?;
-        info!("shader: {}", shader);
-
-        // Create buffers for all outputs of this node
-        for output in current_node.get_output().iter() {
-            if let Some(output_shapes) = shapes_info.get(output) {
-                if output_info
-                    .iter()
-                    .any(|el| el.get_name() == output.as_str())
-                {
-                    buffers.insert(
-                        output.clone(),
-                        resource::buffer(
-                            device,
-                            output_shapes.buffer_len() as _,
-                            output.as_str(),
-                            BufferUsages::STORAGE | BufferUsages::MAP_READ,
-                        ),
-                    );
-                } else {
-                    buffers.insert(
-                        output.clone(),
-                        resource::buffer(
-                            device,
-                            output_shapes.buffer_len() as _,
-                            output.as_str(),
-                            BufferUsages::STORAGE,
-                        ),
-                    );
-                }
-            } else {
-                return Err(OptimizationError::OutputShapeMissing);
-            }
-        }
-
-        // Find all the buffers we want to use as inputs and outputs for this node. These will be 'bound' to variables in the shader
-        let mut binding_counter: usize = 0;
-        let mut entries = vec![];
-
-        for tensor in current_node.get_input() {
-            // Bindings are numbered 0...3 (MAX_BINDINGS_PER_GROUP-1) in binding groups (starting at group 0)
-            let binding_index = (binding_counter % MAX_BINDINGS_PER_GROUP) as u32;
-            entries.push(wgpu::BindGroupEntry {
-                binding: binding_index,
-                resource: buffers
-                    .get(tensor.as_str())
-                    .ok_or_else(|| OptimizationError::TensorMetadataMissing(tensor.to_string()))?
-                    .as_entire_binding(),
-            });
-            binding_counter += 1;
-        }
-
-        for tensor in current_node.get_output() {
-            // Bindings are numbered 0...3 (MAX_BINDINGS_PER_GROUP-1) in binding groups (starting at group 0)
-            let binding_index = (binding_counter % MAX_BINDINGS_PER_GROUP) as u32;
-            entries.push(wgpu::BindGroupEntry {
-                binding: binding_index,
-                resource: buffers
-                    .get(tensor.as_str())
-                    .ok_or_else(|| OptimizationError::TensorMetadataMissing(tensor.to_string()))?
-                    .as_entire_binding(),
-            });
-            binding_counter += 1;
-        }
-
-        // Compile the shader source code for this node
-        let mut bind_groups = vec![];
-        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
-            label: Some(current_node.get_name()),
-            layout: None,
-            module: &device.create_shader_module(&wgpu::ShaderModuleDescriptor {
-                label: Some(current_node.get_name()),
-                source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(&shader)),
-            }),
-            entry_point: "main",
-        });
-
-        // Perform the binding for each group
-        let number_of_groups = ceil(binding_counter as u64, MAX_BINDINGS_PER_GROUP as u64) as usize;
-        for group_index in 0..number_of_groups {
-            let group_range = group_index * MAX_BINDINGS_PER_GROUP
-                ..usize::min(
-                    binding_counter as _,
-                    (group_index + 1) * MAX_BINDINGS_PER_GROUP,
-                );
-            bind_groups.push(device.create_bind_group(&wgpu::BindGroupDescriptor {
-                label: Some(current_node.get_name()),
-                layout: &pipeline.get_bind_group_layout(group_index as u32),
-                entries: &entries[group_range],
-            }));
-        }
-
-        builders.push(EncoderBuilder {
-            pipeline,
-            bind_groups,
-            threads,
-        });
-        node_index += sequence.nodes_consumed;
-    }
-
-    Ok(OptimizedModel { buffers, builders })
-}
diff --git a/src/optimizer.rs b/src/optimizer.rs
new file mode 100644
index 00000000..fb4c4f97
--- /dev/null
+++ b/src/optimizer.rs
@@ -0,0 +1,524 @@
+use protobuf::{ProtobufEnum, RepeatedField};
+use std::{borrow::Cow, collections::HashMap, sync::Arc};
+use thiserror::Error;
+
+use crate::{
+    ir::{Input, Node, NodeDefinition, NodeIdentifier, OperatorDefinition},
+    onnx::TensorProto_DataType,
+    resource::padding,
+    utils::{attribute, get_attribute, AttributeNotFoundError},
+};
+
+#[derive(Debug, Error)]
+pub enum OptimizerError {
+    #[error("node has no inputs")]
+    NoInputs,
+
+    #[error("unsupported: {0}")]
+    Unsupported(String),
+
+    #[error("invalid data type {data_type:?} for input {input} of op {op}")]
+    InvalidDataType {
+        data_type: TensorProto_DataType,
+        input: String,
+        op: String,
+    },
+
+    #[error("required attribute not found: {0}")]
+    AttributeNotFound(#[from] AttributeNotFoundError),
+}
+
+#[derive(Clone)]
+struct Sequence<'model> {
+    node: Arc<Node<'model>>,
+    skip: usize,
+}
+
+pub struct Optimizer<'model> {
+    padded_tensors: HashMap<String, Arc<Node<'model>>>,
+    optimized: HashMap<NodeIdentifier<'model>, Sequence<'model>>,
+}
+
+impl<'model> Optimizer<'model> {
+    pub fn new() -> Self {
+        Self {
+            padded_tensors: HashMap::new(),
+            optimized: HashMap::new(),
+        }
+    }
+
+    /// Optimize a graph
+    pub fn optimize(
+        &mut self,
+        node: Arc<Node<'model>>,
+    ) -> Result<Arc<Node<'model>>, OptimizerError> {
+        let mut path = vec![];
+        let seq = self.optimize_branch_cached(node, &mut path)?;
+        if seq.skip != 0 {
+            panic!("sequencing gone wrong");
+        }
+        Ok(seq.node)
+    }
+
+    fn optimize_branch_cached(
+        &mut self,
+        node: Arc<Node<'model>>,
+        chain: &mut Vec<(String, Arc<Node<'model>>)>,
+    ) -> Result<Sequence<'model>, OptimizerError> {
+        if let Some(optimized) = self.optimized.get(&node.identifier()) {
+            Ok(optimized.clone())
+        } else {
+            let identifier = node.identifier();
+            let optimized = self.optimize_branch(node, chain)?;
+            self.optimized.insert(identifier, optimized.clone());
+            Ok(optimized)
+        }
+    }
+
+    /// Optimize a branch of a graph. This function is recursively called on inputs. The 'chain' parameter is a list of
+    /// operations that follow the node on which optimize_branch is called, and that have exactly one 'dynamic' input
+    /// (either inference input or output of the previous operator). For instance, if the model is A->B->C, then the
+    /// following calls will happen:
+    ///
+    /// - C.optimize_branch(chain = [])
+    ///   - B.optimize_branch(chain = ["C"])
+    ///     - A.optimize_branch(chain = ["C", "B"])
+    ///
+    /// If the model were A -> B+X -> C instead, then the chain would be 'broken' as 'B' has more than one input. The
+    /// 'chain' information can be used to fuse multiple steps. When optimize_branch detects that a chain can be fused,
+    /// it will generate the fused operation and return a Sequence struct with 'skip' set to a positive number.
+    /// This number equals the number of operations that have been fused, and  which  should be omitted from the DAG.
+    /// In the above example, suppose B->C can be fused. This can be detected in the last call:
+    ///
+    /// - C.optimize_branch(chain = []) receives BC(skip=1) from B.optimize_branch, will omit C and return BC(skip=0)
+    ///   - B.optimize_branch(chain = ["C"]) receives BC(skip=2) from A.optimize_branch, will omit B and return BC(skip=1)
+    ///     - A.optimize_branch(chain = ["C", "B"]) detects B->C fusable, returns BC node with skip=2
+    fn optimize_branch(
+        &mut self,
+        node: Arc<Node<'model>>,
+        chain: &mut Vec<(String, Arc<Node<'model>>)>,
+    ) -> Result<Sequence<'model>, OptimizerError> {
+        log::info!(
+            "Optimize {:?} chain length={:?}",
+            node.definition,
+            chain.len()
+        );
+
+        if let NodeDefinition::Operator(op_def) = &node.definition {
+            // Specific operators can be spliced out of the DAG
+            match op_def.proto.get_op_type() {
+                // Identity: we just remove these nodes and connect the input of the destination node to our input's output
+                // A -> Identity -> B => A -> B
+                "Identity" => {
+                    if node.inputs.len() != 1 {
+                        return Err(OptimizerError::NoInputs);
+                    }
+                    return self.optimize_branch_cached(node.inputs[0].source_node.clone(), chain);
+                }
+                // The Dropout operation does nothing when its training_mode is set to 0. We do not support training_mode=1
+                "Dropout" => {
+                    if node.inputs.is_empty() {
+                        return Err(OptimizerError::NoInputs);
+                    }
+                    let training_mode =
+                        get_attribute("training_mode", Some(0), &op_def.proto)? == 1;
+
+                    if training_mode {
+                        return Err(OptimizerError::Unsupported(String::from(
+                            "Dropout with training_mode=1",
+                        )));
+                    }
+
+                    return self.optimize_branch_cached(node.inputs[0].source_node.clone(), chain);
+                }
+                _ => {}
+            }
+
+            // Optimize chains of operators. A chain can only be formed for operations that each have at most one 'dynamic'
+            // input (e.g. inference input or the input from another operation)
+            let dynamic_input_count = node
+                .inputs
+                .iter()
+                .filter(|input| {
+                    matches!(
+                        input.source_node.definition,
+                        NodeDefinition::Operator(..) | NodeDefinition::Input(..)
+                    )
+                })
+                .count();
+
+            if dynamic_input_count == 1 {
+                chain.push((op_def.proto.get_op_type().to_string(), node.clone()));
+
+                if let Some(seq) = self.optimize_chain(chain)? {
+                    log::info!(
+                        "chain optimization: fuse {} operators to {:?}",
+                        seq.skip,
+                        seq.node.definition
+                    );
+                    return Ok(seq);
+                }
+
+                // Continue optimizing inputs, making the chain longer for one input
+                let mut new_inputs = Vec::with_capacity(node.inputs.len());
+
+                for input in &node.inputs {
+                    let source_sequence =
+                        self.optimize_branch_cached(input.source_node.clone(), chain)?;
+
+                    if source_sequence.skip > 0 {
+                        log::info!("operator is optimized away: {:?}", node.definition);
+                        return Ok(Sequence {
+                            node: source_sequence.node,
+                            skip: source_sequence.skip - 1,
+                        });
+                    }
+
+                    new_inputs.push(Input {
+                        source_node: source_sequence.node,
+                        output_index: input.output_index,
+                    });
+                }
+
+                return Ok(Sequence {
+                    node: self.optimized_with(&node, new_inputs)?,
+                    skip: 0,
+                });
+            }
+        }
+
+        // This node has multiple inputs; optimize each input as a new chain
+        let new_inputs = node
+            .inputs
+            .iter()
+            .map(|input| {
+                let mut input_chain = vec![];
+                let source_sequence =
+                    self.optimize_branch_cached(input.source_node.clone(), &mut input_chain)?;
+
+                assert!(
+                    source_sequence.skip == 0,
+                    "cannot skip items in a chain that has multiple inputs for a single op"
+                );
+
+                Ok(Input {
+                    source_node: source_sequence.node,
+                    output_index: input.output_index,
+                })
+            })
+            .collect::<Result<Vec<Input>, OptimizerError>>()?;
+
+        Ok(Sequence {
+            node: self.optimized_with(&node, new_inputs)?,
+            skip: 0,
+        })
+    }
+
+    /// Create a new node from an existing definition, applying optimizations local to a single node
+    fn optimized_with(
+        &mut self,
+        node: &Arc<Node<'model>>,
+        mut new_inputs: Vec<Input<'model>>,
+    ) -> Result<Arc<Node<'model>>, OptimizerError> {
+        match &node.definition {
+            NodeDefinition::Operator(op_def) => {
+                match op_def.proto.get_op_type() {
+                    "Conv" | "ConvRelu" | "ConvLeakyRelu" => {
+                        // This optimization inserts some padding to convolution between kernels with kernel 3x3, because of
+                        // the stride of matrix3x3 is 16 in wgsl. It makes the computation matrixable and increases the performance.
+                        if new_inputs.len() > 2
+                            && get_attribute::<Vec<i64>>("kernel_shape", None, &op_def.proto)?
+                                == [3, 3]
+                            && (get_attribute("pads", Some(vec![0, 0, 0, 0]), &op_def.proto)?
+                                == [1, 1, 1, 1]
+                                || get_attribute(
+                                    "auto_pad",
+                                    Some("SAME_UPPER".to_string()),
+                                    &op_def.proto,
+                                )? == "SAME_UPPER")
+                            && get_attribute("strides", Some(vec![1, 1]), &op_def.proto)? == [1, 1]
+                        {
+                            if let NodeDefinition::Tensor(tensor) =
+                                &new_inputs[1].source_node.definition
+                            {
+                                new_inputs[1] = Input {
+                                    output_index: 0,
+                                    source_node: match self.padded_tensors.get(tensor.get_name()) {
+                                        Some(padded_tensor_node) => padded_tensor_node.clone(),
+                                        None => {
+                                            let data = tensor.get_float_data();
+                                            let raw_data = if !data.is_empty() {
+                                                bytemuck::cast_slice(data)
+                                            } else {
+                                                tensor.get_raw_data()
+                                            };
+
+                                            let padded_raw_data = padding(raw_data, 12, 4);
+
+                                            log::info!(
+                                                "applying padding optimization to tensor {}: strides data is {} bytes before, {} bytes after",
+                                                tensor.get_name(),
+                                                raw_data.len(),
+                                                padded_raw_data.len()
+                                            );
+
+                                            // Create a new tensor with the padded data
+                                            let mut new_tensor = tensor.clone().into_owned();
+                                            new_tensor.set_float_data(vec![]);
+                                            new_tensor.set_raw_data(padded_raw_data);
+                                            let new_node = Arc::new(Node {
+                                                definition: NodeDefinition::Tensor(Box::new(
+                                                    Cow::Owned(new_tensor),
+                                                )),
+                                                inputs: vec![],
+                                            });
+                                            self.padded_tensors.insert(
+                                                tensor.get_name().to_string(),
+                                                new_node.clone(),
+                                            );
+                                            new_node
+                                        }
+                                    },
+                                }
+                            }
+                        }
+
+                        let new_node = Node {
+                            inputs: new_inputs,
+                            definition: NodeDefinition::Operator(op_def.clone()),
+                        };
+
+                        Ok(Arc::new(new_node))
+                    }
+
+                    // The Clip, Split, Resize and Reshape operator each take optional inputs that influence the operation.
+                    // These are typically statically initialized tensors containing shapes. For more efficient execution we
+                    // move these static values to attributes.
+                    op @ ("Clip" | "Split" | "Resize" | "Reshape") => {
+                        if new_inputs.is_empty() {
+                            return Err(OptimizerError::NoInputs);
+                        }
+
+                        // Names of the inputs (see ONNX operator spec)
+                        let attr_names = match op {
+                            "Split" => SPLIT_INPUT_NAMES,
+                            "Resize" => RESIZE_INPUT_NAMES,
+                            "Reshape" => RESHAPE_INPUT_NAMES,
+                            "Clip" => CLIP_INPUT_NAMES,
+                            _ => unreachable!(),
+                        };
+
+                        // Make a new copy of the attributes list (we're going to add attributes)
+                        let mut new_proto = op_def.proto.clone().into_owned();
+                        let mut attributes = op_def.proto.get_attribute().to_vec();
+
+                        // Loop over the inputs (skipping the first one - that's going to be the data input)
+                        for input_index in 1..(new_inputs.len().min(attr_names.len())) {
+                            let source_node = &new_inputs[input_index].source_node;
+                            match &source_node.definition {
+                                // If the input is an initializer (Tensor) we can obtain the data from the definition and move it to an attribute
+                                NodeDefinition::Tensor(tensor_proto) => {
+                                    let attr_name = attr_names[input_index];
+                                    let data_type = TensorProto_DataType::from_i32(
+                                        tensor_proto.get_data_type(),
+                                    )
+                                    .unwrap_or(TensorProto_DataType::UNDEFINED);
+
+                                    match (op, attr_name) {
+                                        // Inputs that need to be converted to an i64 attribute
+                                        ("Split", "split")
+                                        | ("Resize", "roi")
+                                        | ("Resize", "sizes")
+                                        | ("Reshape", "shape") => match data_type {
+                                            TensorProto_DataType::INT64
+                                            | TensorProto_DataType::UNDEFINED => {
+                                                log::info!(
+                                                        "transferring input {} for op {} to i64 attribute (initializer data type: {:?})",
+                                                        attr_name,
+                                                        op,
+                                                        data_type
+                                                    );
+                                                let value = tensor_proto.get_int64_data().to_vec();
+                                                attributes.push(attribute(
+                                                    attr_names[input_index],
+                                                    value,
+                                                ));
+                                            }
+                                            _ => {
+                                                return Err(OptimizerError::InvalidDataType {
+                                                    data_type,
+                                                    input: attr_name.to_string(),
+                                                    op: op.to_string(),
+                                                })
+                                            }
+                                        },
+                                        // Inputs that need to be converted to an f32 attribute
+                                        ("Resize", "scales") => match data_type {
+                                            TensorProto_DataType::FLOAT
+                                            | TensorProto_DataType::UNDEFINED => {
+                                                log::info!(
+                                                        "transferring input {} for op {} to f32 attribute (initializer data type: {:?})",
+                                                        attr_name,
+                                                        op,
+                                                        data_type
+                                                    );
+                                                let value: Vec<f32> =
+                                                    tensor_proto.get_float_data().to_vec();
+                                                attributes.push(attribute(
+                                                    attr_names[input_index],
+                                                    value,
+                                                ));
+                                            }
+                                            _ => {
+                                                return Err(OptimizerError::InvalidDataType {
+                                                    data_type,
+                                                    input: attr_name.to_string(),
+                                                    op: op.to_string(),
+                                                })
+                                            }
+                                        },
+                                        _ => {
+                                            // Some other unspecified input that we do not support yet
+                                            return Err(OptimizerError::Unsupported(format!(
+                                                "data_type {} for input {} to op {}",
+                                                tensor_proto.get_data_type(),
+                                                attr_name,
+                                                op
+                                            )));
+                                        }
+                                    }
+                                }
+                                NodeDefinition::Missing => {
+                                    // Just remove it
+                                }
+                                _ => {
+                                    // One of the inputs (except the first) is something other than a tensor (e.g. 'dynamic')
+                                    return Err(OptimizerError::Unsupported(format!(
+                                        "{} operation with dynamic input for {}",
+                                        op, attr_names[input_index]
+                                    )));
+                                }
+                            }
+                        }
+
+                        // Create new node with extra attributes
+                        new_proto.set_attribute(RepeatedField::from(attributes));
+
+                        let new_node = Node {
+                            inputs: vec![new_inputs[0].clone()],
+                            definition: NodeDefinition::Operator(Box::new(OperatorDefinition {
+                                proto: Cow::Owned(new_proto),
+                                output_shapes: op_def.output_shapes.clone(),
+                            })),
+                        };
+
+                        Ok(Arc::new(new_node))
+                    }
+
+                    _ => Ok(Arc::new(Node {
+                        inputs: new_inputs,
+                        definition: NodeDefinition::Operator(op_def.clone()),
+                    })),
+                }
+            }
+            NodeDefinition::Tensor(..) | NodeDefinition::Input(..) => {
+                assert!(
+                    new_inputs.is_empty(),
+                    "non-operator node cannot have inputs"
+                );
+                // No need to do anything with the provided new inputs
+                Ok(node.clone())
+            }
+            &NodeDefinition::Outputs { .. } => Ok(Arc::new(Node {
+                inputs: new_inputs,
+                definition: node.definition().clone(),
+            })),
+            NodeDefinition::Missing => Ok(node.clone()),
+        }
+    }
+
+    /// Attempt to fuse several operators in a chain of operators with no other dynamic inputs.
+    fn optimize_chain(
+        &mut self,
+        chain: &[(String, Arc<Node<'model>>)],
+    ) -> Result<Option<Sequence<'model>>, OptimizerError> {
+        let path_slices: Vec<&str> = chain.iter().rev().map(|x| x.0.as_str()).collect();
+        match &path_slices[..] {
+            // Conv+Relu or Conv+LeakyRelu: combine into ConvRelu/ConvLeakyRelu
+            ["Conv", "Relu", ..] | ["Conv", "LeakyRelu", ..] => {
+                let conv = chain[chain.len() - 1].1.clone();
+                let relu = chain[chain.len() - 2].1.clone();
+
+                if let (NodeDefinition::Operator(conv_def), NodeDefinition::Operator(relu_def)) =
+                    (&conv.definition, &relu.definition)
+                {
+                    // Use the Conv node as template for the new fused Conv[Leaky]Relu node
+                    let mut convrelu_def = *conv_def.clone();
+                    let mut convrelu_proto = conv_def.proto.clone().into_owned();
+                    let new_op_type = match relu_def.proto.get_op_type() {
+                        "LeakyRelu" => "ConvLeakyRelu",
+                        "Relu" => "ConvRelu",
+                        _ => unreachable!(),
+                    };
+                    convrelu_proto.set_op_type(new_op_type.to_string());
+
+                    // Copy all Relu attributes over to the copy of the Conv node
+                    let mut attributes = conv_def.proto.get_attribute().to_vec();
+                    attributes.extend(relu_def.proto.get_attribute().iter().cloned());
+                    convrelu_proto.set_attribute(RepeatedField::from(attributes));
+                    convrelu_proto.set_name(format!(
+                        "{}+{}",
+                        conv.definition.get_name(),
+                        relu.definition.get_name()
+                    ));
+
+                    log::debug!(
+                        "can fuse chain of Conv/[Leaky]Relu to Conv[Leaky]Relu: {:?}: {:?} + {:?} = {}",
+                        path_slices,
+                        conv.definition(),
+                        relu.definition(),
+                        convrelu_proto.get_name()
+                    );
+
+                    convrelu_def.proto = Cow::Owned(convrelu_proto);
+
+                    let new_inputs = conv
+                        .inputs
+                        .iter()
+                        .map(|input| -> Result<Input, OptimizerError> {
+                            Ok(Input {
+                                source_node: self.optimize(input.source_node.clone())?,
+                                output_index: input.output_index,
+                            })
+                        })
+                        .collect::<Result<Vec<_>, _>>()?;
+
+                    let node = Arc::new(Node {
+                        inputs: conv.inputs.clone(),
+                        definition: NodeDefinition::Operator(Box::new(convrelu_def)),
+                    });
+
+                    Ok(Some(Sequence {
+                        node: self.optimized_with(&node, new_inputs)?,
+                        skip: 1,
+                    }))
+                } else {
+                    unreachable!();
+                }
+            }
+            _ => Ok(None),
+        }
+    }
+}
+
+impl<'model> Default for Optimizer<'model> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+// Names associated with the inputs of the Split, Resize, Reshape and Clip operators (in positional order - see ONNX spec)
+static SPLIT_INPUT_NAMES: &[&str] = &["input", "split"];
+static RESIZE_INPUT_NAMES: &[&str] = &["X", "roi", "scales", "sizes"];
+static RESHAPE_INPUT_NAMES: &[&str] = &["data", "shape"];
+static CLIP_INPUT_NAMES: &[&str] = &["input", "min", "max"];
diff --git a/src/resource.rs b/src/resource.rs
index 9c5a5194..276b29a7 100644
--- a/src/resource.rs
+++ b/src/resource.rs
@@ -10,7 +10,7 @@ pub async fn request_device_queue() -> (wgpu::Device, wgpu::Queue) {
     let adapter = instance
         .request_adapter(&wgpu::RequestAdapterOptionsBase::default())
         .await
-        .expect("No GPU Found for referenced preference");
+        .expect("No GPU found for referenced preference");
 
     // `request_device` instantiates the feature specific connection to the GPU, defining some parameters,
     //  `features` being the available features.
diff --git a/src/sequencer.rs b/src/sequencer.rs
deleted file mode 100644
index 28529474..00000000
--- a/src/sequencer.rs
+++ /dev/null
@@ -1,230 +0,0 @@
-use std::collections::HashMap;
-
-use bytemuck::cast_slice;
-use protobuf::RepeatedField;
-use thiserror::Error;
-use wgpu::{Buffer, BufferUsages, Device};
-
-use crate::{
-    onnx::NodeProto,
-    resource::{self, padding},
-    utils::{attribute, get_attribute, node, AttributeNotFoundError, Shape},
-};
-
-#[derive(Error, Debug)]
-pub enum SequenceError {
-    #[error("a required attribute was not found: {0}")]
-    AttributeNotFound(#[from] AttributeNotFoundError),
-
-    #[error("{0} is not implemented yet")]
-    NotImplemented(String),
-
-    #[error("could not cull identity node")]
-    CullFailed,
-}
-
-pub struct Sequence {
-    pub node: NodeProto,
-    pub nodes_consumed: usize,
-}
-
-/* This function will take in a list of nodes and their op names, and will consume one or more nodes from the front of
-the list and return a single node representing the removed node(s), as well as the number of nodes consumed from the
-sequence. This function can coalesce or even skip certain nodes. */
-pub fn sequence(
-    names: &[&str],
-    nodes: &[NodeProto],
-    device: &Device,
-    initializers: &HashMap<String, &[u8]>,
-    inner_infos: &mut HashMap<String, Buffer>,
-    _shapes_info: &mut HashMap<String, Shape>,
-) -> Result<Sequence, SequenceError> {
-    assert_eq!(names.len(), nodes.len());
-    let inputs = nodes[0].get_input();
-
-    Ok(match names {
-        ["Conv", "Exp", "Add", "Log", "Tanh", "Mul", ..] => Sequence {
-            node: node(
-                inputs.iter().map(|x| x.as_str()).collect(),
-                nodes[6].get_output().iter().map(|x| x.as_str()).collect(),
-                &(nodes[0].get_name().to_string() + nodes[1].get_name()),
-                "ConvMish",
-                nodes[0].get_attribute().to_vec(),
-            ),
-            nodes_consumed: 6,
-        },
-        ["Conv", "Relu", ..] | ["Conv", "LeakyRelu", ..] => {
-            for input in inputs {
-                if let Some(data) = initializers.get(input) {
-                    let data = if input == &inputs[1]
-                        && get_attribute::<Vec<i64>>("kernel_shape", None, &nodes[0])? == [3, 3]
-                        && (get_attribute("pads", Some(vec![0, 0, 0, 0]), &nodes[0])?
-                            == [1, 1, 1, 1]
-                            || get_attribute(
-                                "auto_pad",
-                                Some("SAME_UPPER".to_string()),
-                                &nodes[0],
-                            )? == "SAME_UPPER")
-                        && get_attribute("strides", Some(vec![1, 1]), &nodes[0])? == [1, 1]
-                    {
-                        padding(data, 12, 4)
-                        // data.to_vec()
-                    } else {
-                        data.to_vec()
-                    };
-                    let data = data.to_vec();
-
-                    // debug_assert!(!data.is_empty(), "Not inserting input: {}", input);
-
-                    inner_infos.insert(
-                        input.to_string(),
-                        resource::create_buffer_init(
-                            device,
-                            data.as_slice(),
-                            input,
-                            BufferUsages::STORAGE,
-                        ),
-                    );
-                }
-            }
-
-            let mut attributes = nodes[0].get_attribute().to_vec();
-            for attribute in nodes[1].get_attribute() {
-                attributes.push(attribute.clone());
-            }
-
-            Sequence {
-                nodes_consumed: 2,
-                node: node(
-                    inputs.iter().map(|x| x.as_str()).collect(),
-                    nodes[1].get_output().iter().map(|x| x.as_str()).collect(),
-                    &(nodes[0].get_name().to_string() + nodes[1].get_name()),
-                    "ConvRelu",
-                    attributes,
-                ),
-            }
-        }
-
-        op @ (["Clip", ..] | ["Split", ..] | ["Resize", ..] | ["Reshape", ..]) => {
-            // Remove non binding related input for those Op
-            let mut inputs = inputs.iter();
-
-            // Remove the first input.
-            let input = inputs.next().unwrap();
-            if let Some(data) = initializers.get(input) {
-                inner_infos.insert(
-                    input.to_string(),
-                    resource::create_buffer_init(device, data, input, BufferUsages::STORAGE),
-                );
-            }
-
-            let mut node = nodes[0].clone();
-            node.set_input(RepeatedField::from(vec![input.clone()]));
-            // Transform some intput into attributes for optimisation.
-            let mut attributes = node.take_attribute();
-
-            for input in inputs {
-                match (op, input.as_str()) {
-                    (["Split", ..], "split") => {
-                        let value: Vec<i64> = cast_slice(initializers.get(input).unwrap()).to_vec();
-                        attributes.push(attribute(input, value));
-                    }
-                    (["Resize", ..], "roi") => {
-                        let value: Vec<i64> = cast_slice(initializers.get(input).unwrap()).to_vec();
-                        attributes.push(attribute(input, value));
-                    }
-                    (["Resize", ..], "scales") => {
-                        let value: Vec<f32> = cast_slice(initializers.get(input).unwrap()).to_vec();
-                        attributes.push(attribute(input, value));
-                    }
-                    (["Reshape", ..], "shape") => {
-                        let value: Vec<f32> = cast_slice(initializers.get(input).unwrap()).to_vec();
-                        attributes.push(attribute(input, value));
-                    }
-                    (["Resize", ..], "sizes") => {
-                        let value: Vec<i64> = cast_slice(initializers.get(input).unwrap()).to_vec();
-                        attributes.push(attribute(input, value));
-                    }
-                    _ => (),
-                }
-            }
-
-            node.set_attribute(attributes);
-
-            Sequence {
-                node,
-                nodes_consumed: 1,
-            }
-        }
-        op @ (["Mul", ..] | ["Add", ..]) => {
-            let mut ending_input = vec![];
-            let mut attributes = vec![];
-            for input in inputs {
-                if let Some(data) = initializers.get(input) {
-                    match (data.len(), op) {
-                        (4, ["Mul", ..]) => {
-                            let coeff: Vec<f32> = bytemuck::cast_slice(data).to_vec();
-                            attributes.push(attribute("coefficient", coeff[0]));
-                        }
-                        (12, ["Add", ..]) => {
-                            if data == &[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] {
-                                attributes.push(attribute("coefficient", 0));
-                            } else {
-                                return Err(SequenceError::NotImplemented(String::from(
-                                    "Add with non-zero data",
-                                )));
-                            }
-                        }
-                        _ => {
-                            inner_infos.insert(
-                                input.to_string(),
-                                resource::create_buffer_init(
-                                    device,
-                                    data,
-                                    input,
-                                    BufferUsages::STORAGE,
-                                ),
-                            );
-                        }
-                    }
-                } else {
-                    ending_input.push(input.clone());
-                }
-            }
-
-            let mut node = nodes[0].clone();
-            node.set_input(RepeatedField::from(ending_input));
-            Sequence {
-                node,
-                nodes_consumed: 1,
-            }
-        }
-        [..] => {
-            for input in inputs {
-                if let Some(data) = initializers.get(input) {
-                    // debug_assert!(!data.is_empty(), "Not inserting input: {}", input);
-                    let mut data = data.to_vec();
-
-                    // Prevent issue with minimum buffer size for binding enforced by wgpu
-                    if data.len() < 16 {
-                        data.resize(16, 0);
-                    }
-                    inner_infos.insert(
-                        input.to_string(),
-                        resource::create_buffer_init(
-                            device,
-                            data.as_slice(),
-                            input,
-                            BufferUsages::STORAGE,
-                        ),
-                    );
-                }
-            }
-
-            Sequence {
-                node: nodes[0].clone(),
-                nodes_consumed: 1,
-            }
-        }
-    })
-}
diff --git a/src/utils.rs b/src/utils.rs
index 7de893dd..d6b34022 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -4,7 +4,6 @@ use serde_derive::Serialize;
 use crate::onnx;
 use crate::onnx::OperatorSetIdProto;
 use crate::onnx::ValueInfoProto;
-use std::collections::HashMap;
 use std::convert::From;
 use std::convert::Into;
 use std::fmt::Display;
@@ -13,12 +12,12 @@ use thiserror::Error;
 
 /* Minimum size of a buffer you can create with wgpu. Creating buffers smaller than this leads to panic "Validation
 * error: buffer binding size X is less than minimum 64" in Device::create_bind_group */
-const MINIMUM_BUFFER_SIZE_BYTES: u64 = 64;
+pub const MINIMUM_BUFFER_SIZE_BYTES: u64 = 64;
 
 #[derive(Debug, Serialize, Clone)]
 #[serde(transparent)]
 pub struct Shape {
-    dims: Vec<u64>,
+    pub dims: Vec<u64>,
 }
 
 impl Shape {
@@ -158,15 +157,6 @@ pub fn ceil(num: u64, div: u64) -> u64 {
     num / div + (num % div != 0) as u64
 }
 
-pub fn rename_attribute(
-    attribute: &onnx::AttributeProto,
-    new_name: String,
-) -> onnx::AttributeProto {
-    let mut attr = attribute.clone();
-    attr.set_name(new_name);
-    attr
-}
-
 impl ValueInfoProto {
     pub fn get_shape(&self) -> Shape {
         Shape::from(
@@ -182,48 +172,6 @@ impl ValueInfoProto {
     }
 }
 
-pub fn dimensions_infos(graph_proto: &onnx::GraphProto) -> HashMap<String, Shape> {
-    let mut shapes_info = HashMap::new();
-
-    for info in graph_proto.get_input() {
-        let shape = info.get_shape();
-        shapes_info.insert(info.get_name().to_string(), shape);
-    }
-
-    for info in graph_proto.get_output() {
-        let shape = info.get_shape();
-        shapes_info.insert(info.get_name().to_string(), shape);
-    }
-
-    for info in graph_proto.get_value_info() {
-        let shape = info.get_shape();
-        shapes_info.insert(info.get_name().to_string(), shape);
-    }
-
-    for info in graph_proto.get_initializer() {
-        let shape = Shape::from(info.get_dims());
-        shapes_info.insert(info.get_name().to_string(), shape);
-    }
-
-    shapes_info
-}
-
-pub fn initializers(graph_proto: &onnx::GraphProto) -> HashMap<String, &[u8]> {
-    let mut initializers = HashMap::new();
-    for initializer in graph_proto.get_initializer() {
-        let input = initializer.get_name().to_string();
-        let data = initializer.get_float_data();
-        let raw_data = if !data.is_empty() {
-            bytemuck::cast_slice(data)
-        } else {
-            initializer.get_raw_data()
-        };
-
-        initializers.insert(input, raw_data);
-    }
-    initializers
-}
-
 // TODO: Make dimension optional
 pub fn tensor(name: &str, dimensions: &[i64]) -> onnx::ValueInfoProto {
     let mut dim_value = vec![];
@@ -434,7 +382,7 @@ mod tests {
         let session = pollster::block_on(crate::Session::from_model(conv_model))
             .expect("Session did not create");
 
-        let result = pollster::block_on(session.run(input_data)).unwrap();
+        let result = pollster::block_on(session.run(&input_data)).unwrap();
 
         assert_eq!(
             result["Y"],
diff --git a/templates/endomorphism/activation.wgsl b/templates/endomorphism/activation.wgsl
index 0b78399c..e8247c26 100644
--- a/templates/endomorphism/activation.wgsl
+++ b/templates/endomorphism/activation.wgsl
@@ -1,17 +1,17 @@
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: ArrayVector;
+var<storage, read> input_0: ArrayVector;
 
 [[group(0), binding(1)]]
-var<storage, write> {{ outputs[0] }}: ArrayVector;
+var<storage, write> output_0: ArrayVector;
 
-[[stage(compute), workgroup_size(1)]]
+[[stage(compute), workgroup_size({{ workgroup_size_x }})]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
     let gidx = global_id.x;
 
-{% set activation_input = [ inputs[0], ".data[gidx]"] | join(sep="") %}
-{% set activation_output = [ outputs[0], ".data[gidx]"] | join(sep="") %}
+{% set activation_input = "input_0.data[gidx]" %}
+{% set activation_output = "output_0.data[gidx]" %}
 {% set activation_type = op_type %}
 {%- include "snippets/activation_vec.wgsl" -%}
 
diff --git a/templates/endomorphism/arithmetic.wgsl b/templates/endomorphism/arithmetic.wgsl
index ae4cb36d..d90159f5 100644
--- a/templates/endomorphism/arithmetic.wgsl
+++ b/templates/endomorphism/arithmetic.wgsl
@@ -1,26 +1,29 @@
 {%- include "structs.wgsl" -%}
+
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: ArrayVector;
+var<storage, read> input_0: ArrayVector;
+
+{% if i_lens | length == 2 %}
 
-{% if inputs | length == 2 %}
 [[group(0), binding(1)]]
-var<storage, read> {{ inputs[1] }}: ArrayVector;
+var<storage, read> input_1: ArrayVector;
 
 [[group(0), binding(2)]]
-var<storage, write> {{ outputs[0] }}: ArrayVector;
+var<storage, write> output_0: ArrayVector;
 
 {% else %}
 
 [[group(0), binding(1)]]
-var<storage, write> {{ outputs[0] }}: ArrayVector;
+var<storage, write> output_0: ArrayVector;
 
 {% endif %}
-[[stage(compute), workgroup_size(256, 1, 1)]]
+
+[[stage(compute), workgroup_size({{ workgroup_size_x }}, 1, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
     let gidx = global_id.x;
-{% if inputs | length == 2 %}
-    {{ outputs[0] }}.data[gidx] = {{ inputs[0] }}.data[gidx] {{ op_type }} {{ inputs[1] }}.data[gidx];
+{% if i_lens | length == 2 %}
+    output_0.data[gidx] = input_0.data[gidx] {{ op_type }} input_1.data[gidx];
 {% else %}
-    {{ outputs[0] }}.data[gidx] = {{ inputs[0] }}.data[gidx] {{ op_type }} vec4<f32>({{ coefficient }}., {{ coefficient }}., {{ coefficient }}., {{ coefficient }}.);
+    output_0.data[gidx] = input_0.data[gidx] {{ op_type }} vec4<f32>({{ coefficient }}., {{ coefficient }}., {{ coefficient }}., {{ coefficient }}.);
 {% endif %}
 }
\ No newline at end of file
diff --git a/templates/endomorphism/batchnormalization.wgsl b/templates/endomorphism/batchnormalization.wgsl
index 85a6b4e1..c3b3c299 100644
--- a/templates/endomorphism/batchnormalization.wgsl
+++ b/templates/endomorphism/batchnormalization.wgsl
@@ -7,27 +7,27 @@ struct Block {
 
 // X (input)
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: Block;
+var<storage, read> input_0: Block;
 
 // Scale
 [[group(0), binding(1)]]
-var<storage, read> {{ inputs[1] }}: Array;
+var<storage, read> input_1: Array;
 
 // B (bias)
 [[group(0), binding(2)]]
-var<storage, read> {{ inputs[2] }}: Array;
+var<storage, read> input_2: Array;
 
 // Input mean
 [[group(0), binding(3)]]
-var<storage, read> {{ inputs[3] }}: Array;
+var<storage, read> input_3: Array;
 
 // Input variance
 [[group(1), binding(0)]]
-var<storage, read> {{ inputs[4] }}: Array;
+var<storage, read> input_4: Array;
 
 // Y (Output)
 [[group(1), binding(1)]]
-var<storage, write> {{ outputs[0] }}: Block;
+var<storage, write> output_0: Block;
 
 [[stage(compute), workgroup_size(1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
@@ -36,11 +36,11 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
     let index = global_id.x + batch * {{ batch_size }}u + channel * {{ channel_size }}u;
 
     // Y = (X - input_mean) / sqrt(input_var + epsilon) * scale + B
-    let x = {{ inputs[0] }}.data[index];
-    let channel_scale = {{ inputs[1] }}.data[channel];
-    let channel_bias = {{ inputs[2] }}.data[channel];
-    let channel_mean = {{ inputs[3] }}.data[channel];
-    let channel_var = {{ inputs[4] }}.data[channel];
+    let x = input_0.data[index];
+    let channel_scale = input_1.data[channel];
+    let channel_bias = input_2.data[channel];
+    let channel_mean = input_3.data[channel];
+    let channel_var = input_4.data[channel];
 
-    {{ outputs[0] }}.data[index] = (x - channel_mean) / sqrt(channel_var + {{ epsilon }}) * channel_scale + channel_bias;
+    output_0.data[index] = (x - channel_mean) / sqrt(channel_var + {{ epsilon }}) * channel_scale + channel_bias;
 }
\ No newline at end of file
diff --git a/templates/endomorphism/copy.wgsl b/templates/endomorphism/copy.wgsl
deleted file mode 100644
index 1c82781a..00000000
--- a/templates/endomorphism/copy.wgsl
+++ /dev/null
@@ -1,15 +0,0 @@
-
-{%- include "structs.wgsl" -%}
-
-[[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: ArrayMatrix;
-
-[[group(0), binding(1)]]
-var<storage, write> {{ outputs[0] }}: ArrayMatrix;
-
-
-[[stage(compute), workgroup_size(1)]]
-fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-    let gidx = global_id.x;
-    {{ outputs[0] }}.data[gidx] = {{ inputs[0] }}.data[gidx];
-}
\ No newline at end of file
diff --git a/templates/endomorphism/map.wgsl b/templates/endomorphism/map.wgsl
index 99b46617..c2f12525 100644
--- a/templates/endomorphism/map.wgsl
+++ b/templates/endomorphism/map.wgsl
@@ -1,12 +1,12 @@
 {%- include "structs.wgsl" -%}
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: ArrayVector;
+var<storage, read> input_0: ArrayVector;
 
 [[group(0), binding(1)]]
-var<storage, write> {{ outputs[0] }}: ArrayVector;
+var<storage, write> output_0: ArrayVector;
 
-[[stage(compute), workgroup_size(1)]]
+[[stage(compute), workgroup_size({{ workgroup_size_x }})]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
     let gidx = global_id.x;
-    {{ outputs[0] }}.data[gidx] = {{ op_type | lower }}({{ inputs[0] }}.data[gidx]);
+    output_0.data[gidx] = {{ op_type | lower }}(input_0.data[gidx]);
 }
\ No newline at end of file
diff --git a/templates/endomorphism/softmax.wgsl b/templates/endomorphism/softmax.wgsl
index 7b327226..01fe8fbc 100644
--- a/templates/endomorphism/softmax.wgsl
+++ b/templates/endomorphism/softmax.wgsl
@@ -2,10 +2,10 @@
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: Array;
+var<storage, read> input_0: Array;
 
 [[group(0), binding(1)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 
 [[stage(compute), workgroup_size(1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
@@ -19,21 +19,21 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 	// Therefore we use log(0.0) instead which returns -Infinity
 	var max_element: f32 = log(0.0);
 	for(var k: u32 = 0u; k < {{ i_lens[0] }}u; k = k + 1u) {
-		let element = {{ inputs[0] }}.data[gidx + k];
+		let element = input_0.data[gidx + k];
 		max_element = max(max_element, element);
 	}
 
 	// Calculate sum(exp(input - max(input)))
 	var sum: f32 = 0.0;
 	for(var k: u32 = 0u; k < {{ i_lens[0] }}u; k = k + 1u) {
-		let element = {{ inputs[0] }}.data[gidx + k];
+		let element = input_0.data[gidx + k];
 		sum  = sum + exp(element - max_element);
 	}
 
 	// Calculate elements and write to output
 	for(var k: u32 = 0u; k < {{ i_lens[0] }}u; k = k + 1u) {
-		let element = {{ inputs[0] }}.data[gidx + k];
-		{{ outputs[0] }}.data[gidx + k] = exp(element - max_element) / sum;
+		let element = input_0.data[gidx + k];
+		output_0.data[gidx + k] = exp(element - max_element) / sum;
 	}
 	{% endif %}
 }
\ No newline at end of file
diff --git a/templates/matrix/concat.wgsl b/templates/matrix/concat.wgsl
index cd5a1b59..7fa50e87 100644
--- a/templates/matrix/concat.wgsl
+++ b/templates/matrix/concat.wgsl
@@ -2,27 +2,27 @@
 {%- include "structs.wgsl" -%}
 
 
-{% for input in inputs %}
+{% for input in i_lens %}
 [[group({{ loop.index0 / 4 | int }}), binding({{ loop.index0 % 4}})]]
-var<storage, read> {{ input }}: Array;
+var<storage, read> input_{{ loop.index0 }}: Array;
 {% endfor %}
 
-{% set binding_len = inputs | length %}
+{% set binding_len = i_lens | length %}
 [[group({{ binding_len  / 4 | int }}), binding({{ binding_len % 4 }})]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 
 // concat.wgsl
 [[stage(compute), workgroup_size(256, 1, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
     let gidx = global_id.x;
-{% for input in inputs %}
+{% for input in i_lens %}
 {% if loop.first %}
     if (gidx < {{ i_lens[0] }}u) {
-	    {{ outputs[0] }}.data[gidx] = {{ inputs[0] }}.data[gidx];
+	    output_0.data[gidx] = input_0.data[gidx];
     }
 {% else %}
 	if ((gidx >= {{ cum_len | nth(n=loop.index0 -1) }}u) && (gidx < {{ cum_len | nth(n=loop.index0)}}u)) {
-	    {{ outputs[0] }}.data[gidx] = {{ input }}.data[gidx - {{ cum_len | nth(n=loop.index0 -1) }}u];
+	    output_0.data[gidx] = input_{{ loop.index0 }}.data[gidx - {{ cum_len | nth(n=loop.index0 -1) }}u];
     }
 {% endif %}
 {% endfor %}
diff --git a/templates/matrix/gemm.wgsl b/templates/matrix/gemm.wgsl
index 432953bb..211a4059 100644
--- a/templates/matrix/gemm.wgsl
+++ b/templates/matrix/gemm.wgsl
@@ -1,20 +1,20 @@
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: ArrayVector;
+var<storage, read> input_0: ArrayVector;
 
 [[group(0), binding(1)]]
-var<storage, read> {{ inputs[1] }}: ArrayVector;
+var<storage, read> input_1: ArrayVector;
 
-{%- if inputs | length == 3 -%} // Bias
+{%- if i_lens | length == 3 -%} // Bias
 [[group(0), binding(2)]]
-var<storage, read> {{ inputs[2] }}: ArrayVector;
+var<storage, read> input_2: ArrayVector;
 
 [[group(0), binding(3)]]
-var<storage, write> {{ outputs[0] }}: ArrayVector;
+var<storage, write> output_0: ArrayVector;
 {%- else -%}
 [[group(0), binding(2)]]
-var<storage, write> {{ outputs[0] }}: ArrayVector;
+var<storage, write> output_0: ArrayVector;
 {%- endif -%}  
 
 [[stage(compute), workgroup_size(1)]]
@@ -31,17 +31,17 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
         let index_right = k * {{ i_shape[0][1] }}u + y; 
 
         let mat_left = mat4x4<f32>(
-                              {{ inputs[0] }}.data[index_left], 
-                              {{ inputs[0] }}.data[index_left + {{ i_shape[0][1] / 4 | int }}u],
-                              {{ inputs[0] }}.data[index_left + {{ 2 * i_shape[0][1] / 4 | int }}u],
-                              {{ inputs[0] }}.data[index_left + {{ 3 * i_shape[0][1] / 4 | int }}u],
+                              input_0.data[index_left], 
+                              input_0.data[index_left + {{ i_shape[0][1] / 4 | int }}u],
+                              input_0.data[index_left + {{ 2 * i_shape[0][1] / 4 | int }}u],
+                              input_0.data[index_left + {{ 3 * i_shape[0][1] / 4 | int }}u],
                           );
           
         let mat_right = mat4x4<f32>(
-                              {{ inputs[1] }}.data[index_right], 
-                              {{ inputs[1] }}.data[index_right + {{ i_shape[1][1] / 4 | int }}u],
-                              {{ inputs[1] }}.data[index_right + {{ 2 * i_shape[1][1] / 4 | int }}u],
-                              {{ inputs[1] }}.data[index_right + {{ 3 * i_shape[1][1] / 4 | int }}u],
+                              input_1.data[index_right], 
+                              input_1.data[index_right + {{ i_shape[1][1] / 4 | int }}u],
+                              input_1.data[index_right + {{ 2 * i_shape[1][1] / 4 | int }}u],
+                              input_1.data[index_right + {{ 3 * i_shape[1][1] / 4 | int }}u],
                           );
 	
         product = mat_right * mat_left;
@@ -51,15 +51,15 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 	}
     }
     
-{%- if inputs | length == 3 -%}
-    let bias_row = {{ inputs[2] }}.data[x]; 
+{%- if i_lens | length == 3 -%}
+    let bias_row = input_2.data[x]; 
     var bias = transpose(mat4x4<f32>(bias_row, bias_row, bias_row, bias_row));
     for(var index_mat: u32 = 0u; index_mat < 4u; index_mat = index_mat + 1u) {
-            {{ outputs[0] }}.data[index + index_mat * {{ i_shape[1][1] / 4 | int }}u] = {%- if alpha != 1 -%} {{ alpha | float }} * {%- endif -%}tmpsum[index_mat] + {%- if beta != 1 -%} {{ beta | float }} * {%- endif -%}bias[index_mat];
+            output_0.data[index + index_mat * {{ i_shape[1][1] / 4 | int }}u] = {%- if alpha != 1 -%} {{ alpha | float }} * {%- endif -%}tmpsum[index_mat] + {%- if beta != 1 -%} {{ beta | float }} * {%- endif -%}bias[index_mat];
     }       
 {%- else -%}
     for(var index_mat: u32 = 0u; index_mat < 4u; index_mat = index_mat + 1u) {
-            {{ outputs[0] }}.data[index + index_mat * {{ i_shape[1][1] / 4 | int }}u] = {%- if alpha != 1 -%} {{ alpha | float }} * {%- endif -%}tmpsum[index_mat];
+            output_0.data[index + index_mat * {{ i_shape[1][1] / 4 | int }}u] = {%- if alpha != 1 -%} {{ alpha | float }} * {%- endif -%}tmpsum[index_mat];
     }       
 {%- endif -%}  
 }
\ No newline at end of file
diff --git a/templates/matrix/gemm_1.wgsl b/templates/matrix/gemm_1.wgsl
index b03991e4..7eb91e0f 100644
--- a/templates/matrix/gemm_1.wgsl
+++ b/templates/matrix/gemm_1.wgsl
@@ -1,21 +1,21 @@
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: ArrayVector;
+var<storage, read> input_0: ArrayVector;
 
 [[group(0), binding(1)]]
-var<storage, read> {{ inputs[1] }}: Array;
+var<storage, read> input_1: Array;
 
-{%- if inputs | length == 3 -%} // Bias
+{%- if i_lens | length == 3 -%} // Bias
 [[group(0), binding(2)]]
-var<storage, read> {{ inputs[2] }}: Array;
+var<storage, read> input_2: Array;
 
 [[group(0), binding(3)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 
 {%- else -%}
 [[group(0), binding(2)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 
 {%- endif -%}  
 
@@ -31,13 +31,13 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
         let index_left = k; 
         let index_right = k * {{ i_shape[1][1] * 4 }}u + gidx; 
 
-        let vec_left = {{ inputs[0] }}.data[index_left];
+        let vec_left = input_0.data[index_left];
 
         let vec_right = vec4<f32>(
-                              {{ inputs[1] }}.data[index_right], 
-                              {{ inputs[1] }}.data[index_right + {{ i_shape[1][1] }}u],
-                              {{ inputs[1] }}.data[index_right + {{ 2 * i_shape[1][1] }}u],
-                              {{ inputs[1] }}.data[index_right + {{ 3 * i_shape[1][1] }}u],
+                              input_1.data[index_right], 
+                              input_1.data[index_right + {{ i_shape[1][1] }}u],
+                              input_1.data[index_right + {{ 2 * i_shape[1][1] }}u],
+                              input_1.data[index_right + {{ 3 * i_shape[1][1] }}u],
                           );
 	
         product = dot(vec_left, vec_right);
@@ -45,8 +45,8 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 	    tmpsum = tmpsum + product;
     }
     
-    {{ outputs[0] }}.data[gidx] = {%- if alpha != 1 -%}{{ alpha | float }} * {%- endif -%}tmpsum
-{%- if inputs | length == 3 -%}
- + {%- if beta != 1 -%}{{ beta | float }} * {%- endif -%}{{ inputs[2] }}.data[gidx];
+    output_0.data[gidx] = {%- if alpha != 1 -%}{{ alpha | float }} * {%- endif -%}tmpsum
+{%- if i_lens | length == 3 -%}
+ + {%- if beta != 1 -%}{{ beta | float }} * {%- endif -%}input_2.data[gidx];
 {%- endif -%};
 }
diff --git a/templates/matrix/resize.wgsl b/templates/matrix/resize.wgsl
index de88139e..798fce2b 100644
--- a/templates/matrix/resize.wgsl
+++ b/templates/matrix/resize.wgsl
@@ -2,11 +2,11 @@
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: Array;
+var<storage, read> input_0: Array;
 
 
 [[group(0), binding(1)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 
 // resize.wgsl
 [[stage(compute), workgroup_size(256, 1, 1)]]
@@ -36,6 +36,6 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 			  )) * {{ chunks  }}u 
         {%- endfor -%};
 
-	    {{ outputs[0] }}.data[gidx] = {{ inputs[0] }}.data[index];
+	    output_0.data[gidx] = input_0.data[index];
     }
 }
diff --git a/templates/matrix/split.wgsl b/templates/matrix/split.wgsl
index d77dd357..a40c8257 100644
--- a/templates/matrix/split.wgsl
+++ b/templates/matrix/split.wgsl
@@ -3,12 +3,12 @@
 
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: Array;
+var<storage, read> input_0: Array;
 
 
-{% for output in outputs %}
+{% for output in o_lens %}
 [[group({{ loop.index / 4 | int }}), binding({{ loop.index % 4}})]]
-var<storage, write> {{ output }}: Array;
+var<storage, write> output_{{ loop.index0 }}: Array;
 {% endfor %}
 
 // split.wgsl
@@ -28,7 +28,7 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
         {% endif %}
         {%- endfor -%}
 
-{% for output in outputs %}
+{% for output in o_lens %}
 {%- if loop.first %}
     if (d_{{ axis }} < {{ split | first }}u) {
 
@@ -39,7 +39,7 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
         d_{{ loop.index0 }} * {{ chunk }}u
         {%- endfor -%};
 
-	    {{ output }}.data[index] = {{ inputs[0] }}.data[gidx];
+	    output_{{ loop.index0 }}.data[index] = input_0.data[gidx];
     }
 {%- else %}
 {% set split_output = split | nth(n=loop.index0 -1) %}
@@ -56,7 +56,7 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
         {%- endif -%}
         {%- endfor -%};
 
-	    {{ output }}.data[index] = {{ inputs[0] }}.data[gidx];
+	    output_{{ loop.index0 }}.data[index] = input_0.data[gidx];
     }
 {% endif %}
 {% endfor %}
diff --git a/templates/matrix/transpose.wgsl b/templates/matrix/transpose.wgsl
index 0a07af1c..8725976d 100644
--- a/templates/matrix/transpose.wgsl
+++ b/templates/matrix/transpose.wgsl
@@ -1,10 +1,10 @@
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: Array;
+var<storage, read> input_0: Array;
 
 [[group(0), binding(1)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 
 
 [[stage(compute), workgroup_size(256, 1, 1)]]
@@ -30,6 +30,6 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
         d_{{ loop.index0 }} * {{ perm }}u
         {%- endfor %};
 
-        {{ outputs[0] }}.data[index] = {{ inputs[0] }}.data[gidx];
+        output_0.data[index] = input_0.data[gidx];
     }
 }
\ No newline at end of file
diff --git a/templates/pool/aggregate.wgsl b/templates/pool/aggregate.wgsl
index 44456963..cf4e1e60 100644
--- a/templates/pool/aggregate.wgsl
+++ b/templates/pool/aggregate.wgsl
@@ -1,10 +1,10 @@
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: Array;
+var<storage, read> input_0: Array;
 
 [[group(0), binding(1)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 
 // aggregate.wgsl
 [[stage(compute), workgroup_size(256, 1, 1)]]
@@ -34,10 +34,10 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 
                         let tmp_index  = base_index + tmp_y * {{ original_width }}u + tmp_x;
                         let vector = vec4<f32>(
-                                {{ inputs[0] }}.data[tmp_index],
-                                {{ inputs[0] }}.data[tmp_index + {{ i_chunks[0][1] }}u],
-                                {{ inputs[0] }}.data[tmp_index + {{ 2 * i_chunks[0][1] }}u],
-                                {{ inputs[0] }}.data[tmp_index + {{ 3 * i_chunks[0][1] }}u],
+                                input_0.data[tmp_index],
+                                input_0.data[tmp_index + {{ i_chunks[0][1] }}u],
+                                input_0.data[tmp_index + {{ 2 * i_chunks[0][1] }}u],
+                                input_0.data[tmp_index + {{ 3 * i_chunks[0][1] }}u],
                         );
 {%- if op_type == "MaxPool" -%}
 
@@ -59,7 +59,7 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
         for(var index_vec: u32 = 0u; index_vec < 4u; index_vec = index_vec + 1u) {
                 let index = base_index + index_vec * {{ o_chunks[0][1] }}u;
 
-                {{ outputs[0] }}.data[index] = result[index_vec];
+                output_0.data[index] = result[index_vec];
         }
         }
 }
diff --git a/templates/pool/conv.wgsl b/templates/pool/conv.wgsl
index b592b9dc..82c40883 100644
--- a/templates/pool/conv.wgsl
+++ b/templates/pool/conv.wgsl
@@ -1,21 +1,21 @@
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: Array;
+var<storage, read> input_0: Array;
 
 [[group(0), binding(1)]]
-var<storage, read> {{ inputs[1] }}: Array;
+var<storage, read> input_1: Array;
 
-{%- if inputs | length == 3 -%} // Bias
+{%- if i_lens | length == 3 -%} // Bias
 [[group(0), binding(2)]]
-var<storage, read> {{ inputs[2] }}: Array;
+var<storage, read> input_2: Array;
 
 [[group(0), binding(3)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 
 {%- else -%}
 [[group(0), binding(2)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 {%- endif %}  
 
 // Conv.wgsl
@@ -58,7 +58,7 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
                                 let tmp_index = base_index + tmp_y * {{ original_width }}u + tmp_x;
                                 let index_kernel = base_kernel_index + i * {{ kernel_shape[1] }}u + j;
 
-                                result = {{ inputs[0] }}.data[tmp_index] * {{ inputs[1] }}.data[index_kernel] + result;
+                                result = input_0.data[tmp_index] * input_1.data[index_kernel] + result;
 				
                         }
   	        }
@@ -66,12 +66,12 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
             }
 	}
 
-        {%- if inputs | length == 3 -%}
-        result = result + {{ inputs[2] }}.data[m];
+        {%- if i_lens | length == 3 -%}
+        result = result + input_2.data[m];
         {%- endif -%}
 
 {% set activation_input = "result" %}
-{% set activation_output = [ outputs[0], ".data[gidx]"] | join(sep="") %}
+{% set activation_output = "output_0.data[gidx]" %}
 {% set activation_type = op_type | replace(from="Conv", to="") %}
 {%- include "snippets/activation_scalar.wgsl" -%}
 
diff --git a/templates/pool/conv_kernel_1.wgsl b/templates/pool/conv_kernel_1.wgsl
index c7955a55..5419916b 100644
--- a/templates/pool/conv_kernel_1.wgsl
+++ b/templates/pool/conv_kernel_1.wgsl
@@ -1,20 +1,20 @@
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: Array;
+var<storage, read> input_0: Array;
 
 [[group(0), binding(1)]]
-var<storage, read> {{ inputs[1] }}: ArrayMatrix;
+var<storage, read> input_1: ArrayMatrix;
 
-{%- if inputs | length == 3 -%} // Bias
+{%- if i_lens | length == 3 -%} // Bias
 [[group(0), binding(2)]]
-var<storage, read> {{ inputs[2] }}: ArrayVector;
+var<storage, read> input_2: ArrayVector;
 
 [[group(0), binding(3)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 {%- else -%}
 [[group(0), binding(2)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 {%- endif %}  
 
 // conv_kernel_1.wgsl
@@ -38,19 +38,19 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
                 let base_index = root_index + c * {{ 16 * i_chunks[0][1] }}u;
                 var base_kernel_index = root_kernel_index + c;
                 
-                var matrix_0 = {{ inputs[1] }}.data[base_kernel_index];
-                var matrix_1 = {{ inputs[1] }}.data[base_kernel_index + {{ channel / 16 }}u];
-                var matrix_2 = {{ inputs[1] }}.data[base_kernel_index + {{ 2 * channel / 16 }}u];
-                var matrix_3 = {{ inputs[1] }}.data[base_kernel_index + {{ 3 * channel / 16 }}u];
+                var matrix_0 = input_1.data[base_kernel_index];
+                var matrix_1 = input_1.data[base_kernel_index + {{ channel / 16 }}u];
+                var matrix_2 = input_1.data[base_kernel_index + {{ 2 * channel / 16 }}u];
+                var matrix_3 = input_1.data[base_kernel_index + {{ 3 * channel / 16 }}u];
 
                 for(var index_c_vec: u32 = 0u; index_c_vec < 4u; index_c_vec = index_c_vec + 1u) {
                         
                         let base_index = base_index + index_c_vec * {{ 4 * i_chunks[0][1] }}u;
                         let tmp_vec = vec4<f32>(
-                                {{ inputs[0] }}.data[base_index],
-                                {{ inputs[0] }}.data[base_index + {{ i_chunks[0][1] }}u],
-                                {{ inputs[0] }}.data[base_index + {{ 2 * i_chunks[0][1] }}u],
-                                {{ inputs[0] }}.data[base_index + {{ 3 * i_chunks[0][1] }}u],
+                                input_0.data[base_index],
+                                input_0.data[base_index + {{ i_chunks[0][1] }}u],
+                                input_0.data[base_index + {{ 2 * i_chunks[0][1] }}u],
+                                input_0.data[base_index + {{ 3 * i_chunks[0][1] }}u],
                         );
                 
                         result = tmp_vec * mat4x4<f32>(
@@ -62,8 +62,8 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
                 }
 	}
 
-        {% if inputs | length == 3 -%}
-        result = result + {{ inputs[2] }}.data[m];
+        {% if i_lens | length == 3 -%}
+        result = result + input_2.data[m];
         {%- endif %}
 
 {% set activation_input = "result" %}
@@ -75,7 +75,7 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
         for(var index_vec: u32 = 0u; index_vec < 4u; index_vec = index_vec + 1u) {
                 let index = base_index + index_vec * {{ o_chunks[0][1] }}u;
 
-                {{ outputs[0] }}.data[index] = result[index_vec];
+                output_0.data[index] = result[index_vec];
         }
         }
 }
diff --git a/templates/pool/conv_kernel_3.wgsl b/templates/pool/conv_kernel_3.wgsl
index 05ee2d35..5f4a365c 100644
--- a/templates/pool/conv_kernel_3.wgsl
+++ b/templates/pool/conv_kernel_3.wgsl
@@ -5,21 +5,21 @@ struct ArrayMatrix3 {
 }; // this is used as both input and output for convenience
 
 [[group(0), binding(0)]]
-var<storage, read> {{ inputs[0] }}: Array;
+var<storage, read> input_0: Array;
 
 [[group(0), binding(1)]]
-var<storage, read> {{ inputs[1] }}: ArrayMatrix3;
+var<storage, read> input_1: ArrayMatrix3;
 
-{%- if inputs | length == 3 -%} // Bias
+{%- if i_lens | length == 3 -%} // Bias
 [[group(0), binding(2)]]
-var<storage, read> {{ inputs[2] }}: ArrayVector;
+var<storage, read> input_2: ArrayVector;
 
 [[group(0), binding(3)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 
 {%- else -%}
 [[group(0), binding(2)]]
-var<storage, write> {{ outputs[0] }}: Array;
+var<storage, write> output_0: Array;
 {%- endif %}  
 
 // conv_kernel_3.wgsl
@@ -46,10 +46,10 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
             let base_index = root_index + c * {{ i_chunks[0][1] }}u;
             let base_kernel_index = root_kernel_index + c;
 
-            var kernel_matrix_0 = {{ inputs[1] }}.data[base_kernel_index];
-            var kernel_matrix_1 = {{ inputs[1] }}.data[base_kernel_index + {{ channel }}u];
-            var kernel_matrix_2 = {{ inputs[1] }}.data[base_kernel_index + {{ 2 * channel }}u];
-            var kernel_matrix_3 = {{ inputs[1] }}.data[base_kernel_index + {{ 3 * channel }}u];
+            var kernel_matrix_0 = input_1.data[base_kernel_index];
+            var kernel_matrix_1 = input_1.data[base_kernel_index + {{ channel }}u];
+            var kernel_matrix_2 = input_1.data[base_kernel_index + {{ 2 * channel }}u];
+            var kernel_matrix_3 = input_1.data[base_kernel_index + {{ 3 * channel }}u];
 
             for(var i: u32 = 0u; i < {{ kernel_shape[0] }}u; i = i + 1u) {
 
@@ -68,7 +68,7 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
                                 let tmp_index = base_index + tmp_y * {{ original_width }}u + tmp_x;
                                 let index_kernel = base_kernel_index + i * {{ kernel_shape[1] }}u + j;
                                 
-                                tmp_vec[j] = {{ inputs[0] }}.data[tmp_index];
+                                tmp_vec[j] = input_0.data[tmp_index];
 
                         }
   	        }
@@ -84,8 +84,8 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 
         }       
 	
-        {% if inputs | length == 3 -%}
-        result = result + {{ inputs[2] }}.data[m];
+        {% if i_lens | length == 3 -%}
+        result = result + input_2.data[m];
         {%- endif %}
 
 {% set activation_input = "result" %}
@@ -97,7 +97,7 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
         for(var index_vec: u32 = 0u; index_vec < 4u; index_vec = index_vec + 1u) {
                 let index = base_index + index_vec * {{ o_chunks[0][1] }}u;
 
-                {{ outputs[0] }}.data[index] = result[index_vec];
+                output_0.data[index] = result[index_vec];
         }
         }
 }
diff --git a/templates/snippets/activation_scalar.wgsl b/templates/snippets/activation_scalar.wgsl
index 596c9966..1616a1b7 100644
--- a/templates/snippets/activation_scalar.wgsl
+++ b/templates/snippets/activation_scalar.wgsl
@@ -12,8 +12,8 @@
     {{ activation_output }} = log(1.0 + exp({{ activation_input }}));
 
 {%- elif activation_type == "Clip" -%}
-    let min_clip = {{ inputs[1] }}.data[0u];
-    let max_clip = {{ inputs[2] }}.data[0u];
+    let min_clip = input_1.data[0u];
+    let max_clip = input_2.data[0u];
     {{ activation_output }} = clamp(
         {{ activation_input }}, 
         vec4<f32>(min_clip, min_clip, min_clip, min_clip),
diff --git a/templates/snippets/activation_vec.wgsl b/templates/snippets/activation_vec.wgsl
index 410abdb4..8df53aac 100644
--- a/templates/snippets/activation_vec.wgsl
+++ b/templates/snippets/activation_vec.wgsl
@@ -12,8 +12,8 @@
     {{ activation_output }} = log(vec4<f32>(1.0, 1.0, 1.0, 1.0) + exp({{ activation_input }}));
 
 {%- elif activation_type == "Clip" -%}
-    let min_clip = {{ inputs[1] }}.data[0u];
-    let max_clip = {{ inputs[2] }}.data[0u];
+    let min_clip = input_1.data[0u];
+    let max_clip = input_2.data[0u];
     {{ activation_output }} = clamp(
        {{ activation_input }}, 
         vec4<f32>(min_clip, min_clip, min_clip, min_clip),


Test patch:
diff --git a/tests/batchnormalization.rs b/tests/batchnormalization.rs
index c9007ebb..02d0f7d0 100644
--- a/tests/batchnormalization.rs
+++ b/tests/batchnormalization.rs
@@ -65,7 +65,7 @@ fn batch_normalization() {
     let session =
         pollster::block_on(wonnx::Session::from_model(bn_model)).expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
     let out_y = &result["Y"];
 
     abs_eq_vector(
diff --git a/tests/conv.rs b/tests/conv.rs
index 151788d2..5826ad75 100644
--- a/tests/conv.rs
+++ b/tests/conv.rs
@@ -54,7 +54,7 @@ fn conv_pad() {
     let session =
         pollster::block_on(wonnx::Session::from_model(conv_model)).expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     assert_eq!(
         result["Y"],
@@ -107,7 +107,7 @@ fn conv_without_pad() {
     let session =
         pollster::block_on(wonnx::Session::from_model(conv_model)).expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
     assert_eq!(
         result["Y"],
         [54., 63., 72., 99., 108., 117., 144., 153., 162.]
@@ -160,7 +160,7 @@ fn conv_stride() {
     let session =
         pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     assert_eq!(
         result["Y"],
@@ -209,7 +209,7 @@ fn conv_asymetric_stride() {
     let session =
         pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     assert_eq!(result["Y"], [21., 33., 99., 117., 189., 207., 171., 183.]);
 }
diff --git a/tests/cos.rs b/tests/cos.rs
index bc098715..d219243d 100644
--- a/tests/cos.rs
+++ b/tests/cos.rs
@@ -1,13 +1,8 @@
 use std::collections::HashMap;
-// use wasm_bindgen_test::*;
-// Indicates a f32 overflow in an intermediate Collatz value
-
 use wonnx::utils::{graph, model, node, tensor};
 
 #[test]
 fn test_cos() {
-    // USER INPUT
-
     let n: usize = 16;
     let mut input_data = HashMap::new();
 
@@ -15,7 +10,7 @@ fn test_cos() {
     let shape = vec![n as i64];
     input_data.insert("X".to_string(), data.as_slice());
 
-    // ONNX INPUTS
+    // Model: X -> Cos -> Y
     let model = model(graph(
         vec![tensor("X", &shape)],
         vec![tensor("Y", &shape)],
@@ -24,11 +19,9 @@ fn test_cos() {
         vec![node(vec!["X"], vec!["Y"], "cos", "Cos", vec![])],
     ));
 
-    // LOGIC
-
     let session =
         pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
     assert_eq!(result["Y"], [1.0; 16]);
 }
diff --git a/tests/globalaveragepool.rs b/tests/globalaveragepool.rs
index 6480d156..714fdcec 100644
--- a/tests/globalaveragepool.rs
+++ b/tests/globalaveragepool.rs
@@ -29,6 +29,7 @@ fn global_average_pool() {
     let output_shape = vec![batches as i64, channels as i64, 1, 1];
     input_data.insert("X".to_string(), data.as_slice());
 
+    // Model: X -> GlobalAveragePool -> Y
     let bn_model = model(graph(
         vec![tensor("X", &shape)],
         vec![tensor("Y", &output_shape)],
@@ -43,11 +44,10 @@ fn global_average_pool() {
         )],
     ));
 
-    // LOGIC
     let session =
         pollster::block_on(wonnx::Session::from_model(bn_model)).expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
     let out_y = &result["Y"];
 
     // The GlobalAveragePool op simply averages all pixels in an image (NxCxWxH becomes NxCx1x1). In our test data pixels
diff --git a/tests/identity.rs b/tests/identity.rs
index 779ac5b6..ba15460b 100644
--- a/tests/identity.rs
+++ b/tests/identity.rs
@@ -6,8 +6,6 @@ use wonnx::utils::{graph, model, node, tensor};
 
 #[test]
 fn test_identity() {
-    // USER INPUT
-
     let n: usize = 16;
     let mut input_data = HashMap::new();
 
@@ -15,20 +13,46 @@ fn test_identity() {
     let dims = vec![n as i64];
     input_data.insert("X".to_string(), data.as_slice());
 
-    // ONNX INPUTS
+    // Model: X -> Identity -> Y; Y==Z
     let model = model(graph(
         vec![tensor("X", &dims)],
         vec![tensor("Y", &dims)],
         vec![],
         vec![],
-        vec![node(vec!["X"], vec!["Y"], "cos", "Identity", vec![])],
+        vec![node(vec!["X"], vec!["Y"], "a", "Identity", vec![])],
     ));
 
-    // LOGIC
-
     let session =
         pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
     assert_eq!(result["Y"], data);
 }
+
+#[test]
+fn test_double_identity() {
+    let n: usize = 16;
+    let mut input_data = HashMap::new();
+
+    let data: Vec<f32> = (0..n).map(|x| x as f32).collect();
+    let dims = vec![n as i64];
+    input_data.insert("X".to_string(), data.as_slice());
+
+    // Model: X -> Identity -> Y -> Identity -> Z. X==Z
+    let model = model(graph(
+        vec![tensor("X", &dims)],
+        vec![tensor("Z", &dims)],
+        vec![tensor("Y", &dims)],
+        vec![],
+        vec![
+            node(vec!["X"], vec!["Y"], "a", "Identity", vec![]),
+            node(vec!["Y"], vec!["Z"], "b", "Identity", vec![]),
+        ],
+    ));
+
+    let session =
+        pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
+
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
+    assert_eq!(result["Z"], data);
+}
diff --git a/tests/matrix.rs b/tests/matrix.rs
index bdc415b0..fa6e82f7 100644
--- a/tests/matrix.rs
+++ b/tests/matrix.rs
@@ -1,13 +1,8 @@
-// use approx::assert_relative_eq;
 use std::collections::HashMap;
-// use wasm_bindgen_test::*;
 use wonnx::utils::{attribute, graph, initializer, model, node, tensor};
-// Indicates a f32 overflow in an intermediate Collatz value
 
 #[test]
 fn test_matmul_square_matrix() {
-    // USER INPUT
-
     let n = 16;
     let mut input_data = HashMap::new();
 
@@ -40,19 +35,18 @@ fn test_matmul_square_matrix() {
     let session =
         pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     assert_eq!(result["C"].as_slice(), sum.as_slice().unwrap());
 }
 
 #[test]
 fn test_two_transposes() {
-    // USER INPUT
-
     let mut input_data = HashMap::new();
     let data = (0..2 * 3 * 4).map(|x| x as f32).collect::<Vec<f32>>();
     input_data.insert("X".to_string(), data.as_slice());
 
+    // Model: X -> Transpose -> Y -> Transpose -> Z; X==Z
     let model = model(graph(
         vec![tensor("X", &[2, 3, 4])],
         vec![tensor("Z", &[2, 3, 4])],
@@ -78,15 +72,13 @@ fn test_two_transposes() {
 
     let session =
         pollster::block_on(wonnx::Session::from_model(model)).expect("session did not create");
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     assert_eq!(result["Z"], data);
 }
 
 #[test]
 fn test_split() {
-    // USER INPUT
-
     let mut input_data = HashMap::new();
     let data = (1..=2 * 6).map(|x| x as f32).collect::<Vec<f32>>();
     input_data.insert("X".to_string(), data.as_slice());
@@ -107,7 +99,7 @@ fn test_split() {
 
     let session =
         pollster::block_on(wonnx::Session::from_model(model)).expect("session did not create");
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     let test_y = vec![1., 2., 3., 7., 8., 9.];
     assert_eq!(result["Y"], test_y);
@@ -117,8 +109,7 @@ fn test_split() {
 
 #[test]
 fn test_resize() {
-    // USER INPUT
-
+    let _ = env_logger::builder().is_test(true).try_init();
     let mut input_data = HashMap::new();
     let data = (1..=2 * 4).map(|x| x as f32).collect::<Vec<f32>>();
     input_data.insert("X".to_string(), data.as_slice());
@@ -129,7 +120,7 @@ fn test_resize() {
         vec![],
         vec![initializer("scales", vec![1., 1., 0.6, 0.6])],
         vec![node(
-            vec!["X", "scales"],
+            vec!["X", "" /* roi */, "scales"],
             vec!["Y"],
             "Resize",
             "Resize",
@@ -139,7 +130,7 @@ fn test_resize() {
 
     let session = pollster::block_on(wonnx::Session::from_model(downsampling_model))
         .expect("session did not create");
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     let test_y = vec![1., 3.];
     assert_eq!(result["Y"], test_y);
@@ -164,7 +155,7 @@ fn test_resize() {
 
     let session = pollster::block_on(wonnx::Session::from_model(upsampling_model))
         .expect("session did not create");
-    let _result = pollster::block_on(session.run(input_data)).unwrap();
+    let _result = pollster::block_on(session.run(&input_data)).unwrap();
 
     //let test_y = vec![
     //    1., 1., 1., 2., 2., 2., 1., 1., 1., 2., 2., 2., 3., 3., 3., 4., 4., 4., 3., 3., 3., 4., 4.,
diff --git a/tests/pretrained_models.rs b/tests/pretrained_models.rs
index e596bdc5..24af1f03 100644
--- a/tests/pretrained_models.rs
+++ b/tests/pretrained_models.rs
@@ -15,13 +15,14 @@ fn test_relu() {
         "examples/data/models/single_relu.onnx",
     ))
     .expect("session did not create");
-    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     assert_eq!(result["y"], &[0.0, 1.0]);
 }
 
 #[test]
 fn test_mnist() {
+    let _ = env_logger::builder().is_test(true).try_init();
     let image = load_image("0.jpg");
     let mut input_data = HashMap::new();
     input_data.insert("Input3".to_string(), image.as_slice().unwrap());
@@ -30,7 +31,7 @@ fn test_mnist() {
     ))
     .expect("Session did not create");
 
-    let result = pollster::block_on(session.run(input_data)).unwrap()["Plus214_Output_0"]
+    let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
         .iter()
         .enumerate()
         .fold((0, 0.), |(idx_max, val_max), (idx, val)| {
@@ -46,7 +47,7 @@ fn test_mnist() {
     let image = load_image("3.jpg");
     let mut input_data = HashMap::new();
     input_data.insert("Input3".to_string(), image.as_slice().unwrap());
-    let result = pollster::block_on(session.run(input_data)).unwrap()["Plus214_Output_0"]
+    let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
         .iter()
         .enumerate()
         .fold((0, 0.), |(idx_max, val_max), (idx, val)| {
@@ -62,7 +63,7 @@ fn test_mnist() {
     let image = load_image("5.jpg");
     let mut input_data = HashMap::new();
     input_data.insert("Input3".to_string(), image.as_slice().unwrap());
-    let result = pollster::block_on(session.run(input_data)).unwrap()["Plus214_Output_0"]
+    let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
         .iter()
         .enumerate()
         .fold((0, 0.), |(idx_max, val_max), (idx, val)| {
@@ -78,7 +79,7 @@ fn test_mnist() {
     let image = load_image("7.jpg");
     let mut input_data = HashMap::new();
     input_data.insert("Input3".to_string(), image.as_slice().unwrap());
-    let result = pollster::block_on(session.run(input_data)).unwrap()["Plus214_Output_0"]
+    let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
         .iter()
         .enumerate()
         .fold((0, 0.), |(idx_max, val_max), (idx, val)| {
@@ -92,7 +93,9 @@ fn test_mnist() {
     assert_eq!(result.0, 7);
 }
 
-#[ignore] // Ignore for now because of: https://github.com/gfx-rs/wgpu/issues/2285
+// Ignore on Windows now because of: https://github.com/gfx-rs/wgpu/issues/2285
+// Ignore on Linux because for some reason this does not seem to work properly using Lavapipe
+#[cfg(target_os = "macos")]
 #[test]
 fn test_squeeze() {
     let mut input_data = HashMap::new();
@@ -104,7 +107,7 @@ fn test_squeeze() {
     ))
     .expect("session did not create");
     let result =
-        &pollster::block_on(session.run(input_data)).unwrap()["squeezenet0_flatten0_reshape0"];
+        &pollster::block_on(session.run(&input_data)).unwrap()["squeezenet0_flatten0_reshape0"];
     let mut probabilities = result.iter().enumerate().collect::<Vec<_>>();
 
     probabilities.sort_unstable_by(|a, b| b.1.partial_cmp(a.1).unwrap());