Repository: webonnx/wonnx (language: python)

Issue:
Make wonnx more robust by using a DAG for optimization and transformation
@haixuanTao this is a work-in-progress for adding an 'intermediate representation' - i.e. a tree of nodes that we can more easily perform operations on (such as optimizations). It can also be used to fixed the 'execution order' thing (see #44). 


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index d877ce06..4d531e63 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -7,7 +7,7 @@ on:
 
 env:
   RUST_BACKTRACE: 1
-  RUST_VERSION: 1.55
+  RUST_VERSION: 1.58
 
 # We distinguish the following kinds of builds:
 # - local: build for the same target as we compile on, and do local tests
@@ -33,7 +33,7 @@ jobs:
             target: x86_64-pc-windows-msvc
             tool: check
             kind: other
-          
+
           # MacOS
           - name: MacOS x86_64
             os: macos-10.15
@@ -48,7 +48,7 @@ jobs:
             tool: check
             # Mac has no software runners, so don't run tests
             kind: other
-          
+
           # IOS
           - name: IOS aarch64
             os: macos-10.15
@@ -56,7 +56,6 @@ jobs:
             tool: clippy
             kind: other
 
-
           # Linux
           - name: Linux x86_64
             os: ubuntu-20.04
@@ -77,7 +76,6 @@ jobs:
             tool: check
             kind: other
 
-          
           # Android
           - name: Android aarch64
             os: ubuntu-20.04
@@ -85,7 +83,6 @@ jobs:
             tool: clippy
             kind: other
 
-          
           # WebGPU/WebGL
           - name: WebAssembly
             os: ubuntu-20.04
@@ -151,15 +148,15 @@ jobs:
           mkdir .cargo
           echo """[profile.dev]
           debug = 1" > .cargo/config.toml
-    #  - name: check web
-    #    if: matrix.kind == 'web'
-    #    run: |
-    #      cargo ${{matrix.tool}} --target ${{ matrix.target }} 
+      #  - name: check web
+      #    if: matrix.kind == 'web'
+      #    run: |
+      #      cargo ${{matrix.tool}} --target ${{ matrix.target }}
       - name: check native
         if: matrix.kind == 'local' || matrix.kind == 'other'
         run: |
           # check with no features
-          cargo ${{matrix.tool}} -p wonnx --target ${{ matrix.target }} 
+          cargo ${{matrix.tool}} -p wonnx --target ${{ matrix.target }}
       - name: local tests
         if: matrix.kind == 'local'
         shell: bash
@@ -188,4 +185,4 @@ jobs:
 
       - name: run rustfmt
         run: |
-          cargo fmt -- --check
\ No newline at end of file
+          cargo fmt -- --check
diff --git a/README.md b/README.md
index 51cfe1ab..93cb6da5 100644
--- a/README.md
+++ b/README.md
@@ -83,7 +83,7 @@ cargo run --example mnist --release
 fn main() -> HashMap<String, Vec<f32>> {
     let mut input_data = HashMap::new();
     let image = load_squeezenet_image(); // Load image
-    input_data.insert("data".to_string(), image.as_slice().unwrap());
+    input_data.insert("data".to_string(), InputTensor::F32(image.as_slice().unwrap()));
 
     let session = pollster::block_on(wonnx::Session::from_path(
         "examples/data/models/opt-squeeze.onnx",
@@ -197,7 +197,7 @@ fn test_matmul_square_matrix() {
 |<a href="#AveragePool">AveragePool</a>|<a href="Changelog.md#AveragePool-11">11</a>, <a href="Changelog.md#AveragePool-10">10</a>, <a href="Changelog.md#AveragePool-7">7</a>, <a href="Changelog.md#AveragePool-1">1</a>|✅|
 |<a href="#BatchNormalization">BatchNormalization</a>|<a href="Changelog.md#BatchNormalization-15">15</a>, <a href="Changelog.md#BatchNormalization-14">14</a>, <a href="Changelog.md#BatchNormalization-9">9</a>, <a href="Changelog.md#BatchNormalization-7">7</a>, <a href="Changelog.md#BatchNormalization-6">6</a>, <a href="Changelog.md#BatchNormalization-1">1</a>|✅|
 |<a href="#BitShift">BitShift</a>|<a href="Changelog.md#BitShift-11">11</a>|
-|<a href="#Cast">Cast</a>|<a href="Changelog.md#Cast-13">13</a>, <a href="Changelog.md#Cast-9">9</a>, <a href="Changelog.md#Cast-6">6</a>, <a href="Changelog.md#Cast-1">1</a>|
+|<a href="#Cast">Cast</a>|<a href="Changelog.md#Cast-13">13</a>, <a href="Changelog.md#Cast-9">9</a>, <a href="Changelog.md#Cast-6">6</a>, <a href="Changelog.md#Cast-1">1</a>|✅|
 |<a href="#Ceil">Ceil</a>|<a href="Changelog.md#Ceil-13">13</a>, <a href="Changelog.md#Ceil-6">6</a>, <a href="Changelog.md#Ceil-1">1</a>|✅|
 |<a href="#Clip">Clip</a>|<a href="Changelog.md#Clip-13">13</a>, <a href="Changelog.md#Clip-12">12</a>, <a href="Changelog.md#Clip-11">11</a>, <a href="Changelog.md#Clip-6">6</a>, <a href="Changelog.md#Clip-1">1</a>|✅|
 |<a href="#Compress">Compress</a>|<a href="Changelog.md#Compress-11">11</a>, <a href="Changelog.md#Compress-9">9</a>|
@@ -281,7 +281,7 @@ fn test_matmul_square_matrix() {
 |<a href="#RandomNormalLike">RandomNormalLike</a>|<a href="Changelog.md#RandomNormalLike-1">1</a>|
 |<a href="#RandomUniform">RandomUniform</a>|<a href="Changelog.md#RandomUniform-1">1</a>|
 |<a href="#RandomUniformLike">RandomUniformLike</a>|<a href="Changelog.md#RandomUniformLike-1">1</a>|
-|<a href="#Reciprocal">Reciprocal</a>|<a href="Changelog.md#Reciprocal-13">13</a>, <a href="Changelog.md#Reciprocal-6">6</a>, <a href="Changelog.md#Reciprocal-1">1</a>|
+|<a href="#Reciprocal">Reciprocal</a>|<a href="Changelog.md#Reciprocal-13">13</a>, <a href="Changelog.md#Reciprocal-6">6</a>, <a href="Changelog.md#Reciprocal-1">1</a>|✅|
 |<a href="#ReduceL1">ReduceL1</a>|<a href="Changelog.md#ReduceL1-13">13</a>, <a href="Changelog.md#ReduceL1-11">11</a>, <a href="Changelog.md#ReduceL1-1">1</a>|
 |<a href="#ReduceL2">ReduceL2</a>|<a href="Changelog.md#ReduceL2-13">13</a>, <a href="Changelog.md#ReduceL2-11">11</a>, <a href="Changelog.md#ReduceL2-1">1</a>|
 |<a href="#ReduceLogSum">ReduceLogSum</a>|<a href="Changelog.md#ReduceLogSum-13">13</a>, <a href="Changelog.md#ReduceLogSum-11">11</a>, <a href="Changelog.md#ReduceLogSum-1">1</a>|
diff --git a/py-wonnx/src/lib.rs b/py-wonnx/src/lib.rs
index e1e1a4d9..48e3a259 100644
--- a/py-wonnx/src/lib.rs
+++ b/py-wonnx/src/lib.rs
@@ -2,6 +2,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyDict;
 use std::collections::HashMap;
 
+use wonnx::utils::InputTensor;
 use wonnx::Session;
 #[pyclass]
 #[repr(transparent)]
@@ -27,7 +28,7 @@ impl PySession {
         let map: HashMap<String, Vec<f32>> = dict.extract().unwrap();
         let mut inputs = HashMap::new();
         for (key, value) in map.iter() {
-            inputs.insert(key.clone(), value.as_slice());
+            inputs.insert(key.clone(), InputTensor::F32(value.as_slice()));
         }
         let result = pollster::block_on(self.session.run(&inputs)).unwrap();
 
diff --git a/wonnx/examples/mnist.rs b/wonnx/examples/mnist.rs
index 33f5ab16..9a5f387f 100644
--- a/wonnx/examples/mnist.rs
+++ b/wonnx/examples/mnist.rs
@@ -5,6 +5,7 @@ use std::collections::HashMap;
 use image::{imageops::FilterType, ImageBuffer, Pixel, Rgb};
 use std::path::Path;
 use std::time::Instant;
+use wonnx::utils::InputTensor;
 
 // Args Management
 async fn run() {
@@ -25,11 +26,16 @@ async fn execute_gpu() -> Option<HashMap<String, Vec<f32>>> {
     let mut input_data = HashMap::new();
 
     let image = load_image();
-    input_data.insert("Input3".to_string(), image.as_slice().unwrap());
+    input_data.insert(
+        "Input3".to_string(),
+        InputTensor::F32(image.as_slice().unwrap()),
+    );
+
     let model_path = Path::new(env!("CARGO_MANIFEST_DIR"))
         .join("examples/data/models")
         .join("opt-squeeze.onnx");
     let session = wonnx::Session::from_path(model_path).await.unwrap();
+
     let time_pre_compute = Instant::now();
     let result = session.run(&input_data).await.unwrap();
     let time_post_compute = Instant::now();
diff --git a/wonnx/examples/simple_graph.rs b/wonnx/examples/simple_graph.rs
index a6a873a9..0551e2e8 100644
--- a/wonnx/examples/simple_graph.rs
+++ b/wonnx/examples/simple_graph.rs
@@ -2,7 +2,7 @@ use std::collections::HashMap;
 // use wasm_bindgen_test::*;
 // Indicates a f32 overflow in an intermediate Collatz value
 use wonnx::{
-    utils::{attribute, graph, initializer, model, node, tensor},
+    utils::{attribute, graph, initializer, model, node, tensor, InputTensor},
     SessionError, WonnxError,
 };
 
@@ -22,7 +22,7 @@ async fn execute_gpu() -> Result<HashMap<String, Vec<f32>>, SessionError> {
     let mut input_data = HashMap::new();
 
     let data: Vec<f32> = (0..25).map(|x| x as f32).collect();
-    input_data.insert("X".to_string(), data.as_slice());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // ONNX INPUTS
     let shape = vec![1, c as i64, n as i64, n as i64];
diff --git a/wonnx/examples/squeeze.rs b/wonnx/examples/squeeze.rs
index db8f423c..eccfcb5c 100644
--- a/wonnx/examples/squeeze.rs
+++ b/wonnx/examples/squeeze.rs
@@ -11,6 +11,7 @@ use std::{
     path::Path,
     time::Duration,
 };
+use wonnx::utils::InputTensor;
 use wonnx::WonnxError;
 
 // Args Management
@@ -34,7 +35,10 @@ async fn run() {
 async fn execute_gpu() -> Result<HashMap<String, Vec<f32>>, WonnxError> {
     let mut input_data = HashMap::new();
     let image = load_image();
-    input_data.insert("data".to_string(), image.as_slice().unwrap());
+    input_data.insert(
+        "data".to_string(),
+        InputTensor::F32(image.as_slice().unwrap()),
+    );
 
     let model_path = Path::new(env!("CARGO_MANIFEST_DIR"))
         .join("examples/data/models")
diff --git a/wonnx/src/compiler.rs b/wonnx/src/compiler.rs
index b04dd952..946d58cd 100644
--- a/wonnx/src/compiler.rs
+++ b/wonnx/src/compiler.rs
@@ -1,4 +1,6 @@
-use crate::utils::{ceil, get_attribute, AttributeNotFoundError, DataType, Shape};
+use crate::utils::{
+    ceil, get_attribute, AttributeNotFoundError, DataTypeError, MultiType, ScalarType, Shape,
+};
 use tera::{Context, Tera};
 use thiserror::Error;
 
@@ -40,6 +42,11 @@ lazy_static! {
             include_str!("../templates/endomorphism/map.wgsl"),
         )
         .unwrap();
+        tera.add_raw_template(
+            "endomorphism/cast.wgsl",
+            include_str!("../templates/endomorphism/cast.wgsl"),
+        )
+        .unwrap();
         tera.add_raw_template(
             "matrix/concat.wgsl",
             include_str!("../templates/matrix/concat.wgsl"),
@@ -146,6 +153,55 @@ pub enum CompileError {
 
     #[error("the model exceeds the limit for {0}: {1} > {2}")]
     ComputeLimitExceeded(String, u32, u32),
+
+    #[error("cannot determine data type to use: {0} or {1}")]
+    TypesDisagree(ScalarType, ScalarType),
+
+    #[error("cannot infer data type to use")]
+    TypeUnderspecified,
+
+    #[error("invalid type encountered: {0}")]
+    InvalidType(#[from] DataTypeError),
+}
+
+struct NodeTemplate {
+    scalar_type: ScalarType,
+    template: &'static str,
+    threads: (u32, u32, u32),
+}
+
+/// Returns the data type of the input and output shapes, but error if these types differ or when no input/output was specified
+fn agreed_type(
+    input_shapes: &[&Shape],
+    output_shapes: &[&Shape],
+) -> Result<ScalarType, CompileError> {
+    let mut data_type: Option<ScalarType> = None;
+
+    for input_shape in input_shapes {
+        let input_type = input_shape.data_type;
+        match data_type {
+            Some(dt) => {
+                if dt != input_type {
+                    return Err(CompileError::TypesDisagree(dt, input_type));
+                }
+            }
+            None => data_type = Some(input_type),
+        }
+    }
+
+    for output_shape in output_shapes {
+        let output_type = output_shape.data_type;
+        match data_type {
+            Some(dt) => {
+                if dt != output_type {
+                    return Err(CompileError::TypesDisagree(dt, output_type));
+                }
+            }
+            None => data_type = Some(output_type),
+        }
+    }
+
+    data_type.ok_or(CompileError::TypeUnderspecified)
 }
 
 pub fn compile(
@@ -166,18 +222,20 @@ pub fn compile(
 
     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 i_dims: Vec<&Vec<u64>> = input_shapes.iter().map(|s| &s.dims).collect();
+    let o_dims: Vec<&Vec<u64>> = output_shapes.iter().map(|s| &s.dims).collect();
 
     let mut context = Context::new();
     context.insert("i_lens", &input_lengths);
     context.insert("o_lens", &output_lengths);
-    context.insert("i_shape", &input_shapes);
-    context.insert("o_shape", &output_shapes);
+    context.insert("i_shape", &i_dims);
+    context.insert("o_shape", &o_dims);
     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() {
+    let node_template: NodeTemplate = 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()));
@@ -185,14 +243,36 @@ pub fn compile(
 
         // Map simple function
         "Abs" | "Acos" | "Asin" | "Atan" | "Ceil" | "Cos" | "Cosh" | "Exp" | "Floor" | "Log"
-        | "Round" | "Sign" | "Sin" | "Sinh" | "Sqrt" | "Tan" | "Tanh" => {
+        | "Round" | "Sign" | "Sin" | "Sinh" | "Sqrt" | "Tan" | "Tanh" | "Reciprocal" => {
+            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);
+            NodeTemplate {
+                scalar_type: agreed_type(input_shapes, output_shapes)?,
+                template: "endomorphism/map.wgsl",
+                threads: (x_threads, 1, 1),
+            }
+        }
+
+        "Cast" => {
+            let cast_to_type =
+                ScalarType::from_i32(get_attribute::<i64>("to", None, node)? as i32)?;
+            context.insert("cast_to_type", cast_to_type.wgsl_type_name());
+
             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)
+            NodeTemplate {
+                scalar_type: agreed_type(input_shapes, &[])?,
+                template: "endomorphism/cast.wgsl",
+                threads: (x_threads, 1, 1),
+            }
         }
 
         "Softmax" => {
@@ -237,7 +317,11 @@ pub fn compile(
                 });
             }
 
-            ("endomorphism/softmax.wgsl".to_string(), 1, 1, 1)
+            NodeTemplate {
+                scalar_type: agreed_type(input_shapes, output_shapes)?,
+                template: "endomorphism/softmax.wgsl",
+                threads: (1, 1, 1),
+            }
         }
 
         // Arithmetic operation
@@ -269,13 +353,17 @@ pub fn compile(
             );
 
             let (x_threads, workgroup_size_x) = workgroup_size(
-                ceil(output_lengths[0], 1024) as _,
+                ceil(output_lengths[0], 4) as _,
                 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)
+            NodeTemplate {
+                scalar_type: agreed_type(input_shapes, output_shapes)?,
+                template: "endomorphism/arithmetic.wgsl",
+                threads: (x_threads, 1, 1),
+            }
         }
         // Not taking into account attributes
         "BatchNormalization" => {
@@ -331,10 +419,10 @@ pub fn compile(
             }
 
             // If w*h is a multiple of 4, we can use vec4 in our shader
-            let elem_type = DataType::for_size((input_w * input_h) as usize);
+            let elem_type = MultiType::for_size((input_w * input_h) as usize, ScalarType::F32);
 
             context.insert("elem_type", &elem_type.wgsl_type_name());
-            context.insert("elem_stride", &elem_type.size_bytes());
+            context.insert("elem_stride", &elem_type.stride());
 
             // The default for epsilon is 1e05, see https://github.com/onnx/onnx/blob/master/docs/Changelog.md#attributes-252
             let epsilon = get_attribute("epsilon", Some(1e-05), node)?;
@@ -351,12 +439,15 @@ pub fn compile(
                 &ceil(input_w * input_h, elem_type.elements() as u64),
             );
 
-            (
-                "endomorphism/batchnormalization.wgsl".to_string(),
-                ceil(input_w * input_h, elem_type.elements() as u64) as _,
-                input_channels as _,
-                input_batches as _,
-            )
+            NodeTemplate {
+                scalar_type: agreed_type(&input_shapes[0..1], &output_shapes[0..1])?,
+                template: "endomorphism/batchnormalization.wgsl",
+                threads: (
+                    ceil(input_w * input_h, elem_type.elements() as u64) as _,
+                    input_channels as _,
+                    input_batches as _,
+                ),
+            }
         }
         "Relu" | "Sigmoid" | "Softsign" | "Softplus" | "Clip" | "Celu" | "Elu" | "LeakyRelu" => {
             let alpha = get_attribute("alpha", Some(1.0), node)?;
@@ -370,12 +461,11 @@ pub fn compile(
 
             context.insert("workgroup_size_x", &workgroup_size_x);
 
-            (
-                "endomorphism/activation.wgsl".to_string(),
-                x_threads as _,
-                1,
-                1,
-            )
+            NodeTemplate {
+                scalar_type: agreed_type(input_shapes, output_shapes)?,
+                template: "endomorphism/activation.wgsl",
+                threads: (x_threads, 1, 1),
+            }
         }
         "Concat" => {
             let mut input_cumulative_len = vec![];
@@ -385,12 +475,12 @@ pub fn compile(
                 input_cumulative_len.push(sum);
             }
             context.insert("cum_len", &input_cumulative_len);
-            (
-                "matrix/concat.wgsl".to_string(),
-                ceil(output_lengths[0], 256) as u32,
-                1,
-                1,
-            )
+
+            NodeTemplate {
+                scalar_type: agreed_type(input_shapes, output_shapes)?,
+                template: "matrix/concat.wgsl",
+                threads: (ceil(output_lengths[0], 256) as u32, 1, 1),
+            }
         }
         op @ ("MaxPool" | "AveragePool" | "Conv" | "ConvRelu" | "ConvLeakyRelu" | "ConvMish"
         | "GlobalAveragePool") => {
@@ -475,12 +565,11 @@ pub fn compile(
 
             // GLSL shader for convolution computation
             match op {
-                "MaxPool" | "AveragePool" | "GlobalAveragePool" => (
-                    "pool/aggregate.wgsl".to_string(),
-                    ceil(output_lengths[0], 1024) as _,
-                    1,
-                    1,
-                ),
+                "MaxPool" | "AveragePool" | "GlobalAveragePool" => NodeTemplate {
+                    scalar_type: agreed_type(input_shapes, &output_shapes[0..1])?,
+                    template: "pool/aggregate.wgsl",
+                    threads: (ceil(output_lengths[0], 1024) as _, 1, 1),
+                },
                 "Conv" | "ConvRelu" | "ConvLeakyRelu" | "ConvMish" => {
                     // Alpha is the Leaky Relu attribute
                     let alpha = get_attribute("alpha", Some(0.01), node)?;
@@ -493,30 +582,27 @@ pub fn compile(
                         && (input_shape.dim(1) % 16 == 0)
                         && (output_shape.dim(1) % 4 == 0)
                     {
-                        (
-                            "pool/conv_kernel_1.wgsl".to_string(),
-                            ceil(output_lengths[0], 1024) as _,
-                            1,
-                            1,
-                        )
+                        NodeTemplate {
+                            scalar_type: agreed_type(input_shapes, output_shapes)?,
+                            template: "pool/conv_kernel_1.wgsl",
+                            threads: (ceil(output_lengths[0], 1024) as _, 1, 1),
+                        }
                     } else if (strides == [1, 1])
                         && (kernel_shape == [3, 3])
                         && (dilations == [1, 1])
                         && (output_shape.dim(1) % 4 == 0)
                     {
-                        (
-                            "pool/conv_kernel_3.wgsl".to_string(),
-                            ceil(output_lengths[0], 1024) as _,
-                            1,
-                            1,
-                        )
+                        NodeTemplate {
+                            scalar_type: agreed_type(input_shapes, output_shapes)?,
+                            template: "pool/conv_kernel_3.wgsl",
+                            threads: (ceil(output_lengths[0], 1024) as _, 1, 1),
+                        }
                     } else {
-                        (
-                            "pool/conv.wgsl".to_string(),
-                            ceil(output_lengths[0], 256) as _,
-                            1,
-                            1,
-                        )
+                        NodeTemplate {
+                            scalar_type: agreed_type(input_shapes, output_shapes)?,
+                            template: "pool/conv.wgsl",
+                            threads: (ceil(output_lengths[0], 256) as _, 1, 1),
+                        }
                     }
                 }
                 _ => return Err(CompileError::InvalidOperation(op.to_string())),
@@ -525,6 +611,21 @@ pub fn compile(
         op @ ("Gemm" | "MatMul") => {
             let alpha = get_attribute("alpha", Some(1.0), node)?;
             let beta = get_attribute("beta", Some(1.0), node)?;
+
+            // Whether A resp. B should be transposed, or C should be broadcast (default: 0 = false)
+            if op == "Gemm" {
+                let transpose_a = get_attribute("transA", Some(0), node)?;
+                let transpose_b = get_attribute("transB", Some(0), node)?;
+                let broadcast = get_attribute("broadcast", Some(0), node)?;
+
+                if transpose_a != 0 || transpose_b != 0 || broadcast != 0 {
+                    return Err(CompileError::UnimplementedVariant {
+                        variant: "Gemm with transA/transB/broadcast not equal to zero".to_string(),
+                        op: op.to_string(),
+                    });
+                }
+            }
+
             context.insert("alpha", &alpha);
             context.insert("beta", &beta);
 
@@ -543,11 +644,21 @@ pub fn compile(
             }
 
             if input_shapes[0].dim(0) == 1 {
-                let threads = output_shapes[0].dim(1);
-                ("matrix/gemm_1.wgsl".to_string(), threads as _, 1, 1)
+                NodeTemplate {
+                    scalar_type: agreed_type(input_shapes, output_shapes)?,
+                    template: "matrix/gemm_1.wgsl",
+                    threads: (output_shapes[0].dim(1) as _, 1, 1),
+                }
             } else {
-                let threads = input_shapes[0].dim(0) * input_shapes[1].dim(1) / 16;
-                ("matrix/gemm.wgsl".to_string(), threads as _, 1, 1)
+                NodeTemplate {
+                    scalar_type: agreed_type(input_shapes, output_shapes)?,
+                    template: "matrix/gemm.wgsl",
+                    threads: (
+                        (input_shapes[0].dim(0) * input_shapes[1].dim(1) / 16) as _,
+                        1,
+                        1,
+                    ),
+                }
             }
         }
         "Resize" => {
@@ -640,12 +751,11 @@ pub fn compile(
             let exclude_outside = get_attribute("exclude_outside", Some(0), node)?;
             context.insert("exclude_outside", &exclude_outside);
 
-            (
-                "matrix/resize.wgsl".to_string(),
-                ceil(output_lengths[0], 256) as u32,
-                1,
-                1,
-            )
+            NodeTemplate {
+                scalar_type: agreed_type(&input_shapes[0..1], &output_shapes[0..1])?,
+                template: "matrix/resize.wgsl",
+                threads: (ceil(output_lengths[0], 256) as u32, 1, 1),
+            }
         }
         "Sum" => return Err(CompileError::UnimplementedOp(String::from("Sum"))),
         "Split" => {
@@ -663,12 +773,11 @@ pub fn compile(
             let split = get_attribute::<Vec<i64>>("split", Some(default_split), node)?;
             context.insert("split", &split);
 
-            (
-                "matrix/split.wgsl".to_string(),
-                ceil(output_lengths[0], 256) as u32,
-                1,
-                1,
-            )
+            NodeTemplate {
+                scalar_type: agreed_type(&input_shapes[0..1], &output_shapes[0..1])?,
+                template: "matrix/split.wgsl",
+                threads: (ceil(output_lengths[0], 256) as u32, 1, 1),
+            }
         }
         "Transpose" => {
             let default = ((input_lengths[0] as i64)..0).collect::<Vec<_>>();
@@ -686,46 +795,59 @@ pub fn compile(
 
             context.insert("permuted_chunks", &chunks);
 
-            (
-                "matrix/transpose.wgsl".to_string(),
-                ceil(output_lengths[0], 256) as _,
-                1,
-                1,
-            )
+            NodeTemplate {
+                scalar_type: agreed_type(input_shapes, output_shapes)?,
+                template: "matrix/transpose.wgsl",
+                threads: (ceil(output_lengths[0], 256) as _, 1, 1),
+            }
         }
         op => return Err(CompileError::UnimplementedOp(op.to_string())),
     };
 
-    let shader = TEMPLATES
-        .render(&template, &context)
-        .expect("failed to render shader");
-
     // Check if we remain within the limits of the thread count allowed by WebGPU
-    if x > MAX_COMPUTE_WORKGROUPS_PER_DIMENSION {
+    if node_template.threads.0 > MAX_COMPUTE_WORKGROUPS_PER_DIMENSION {
         return Err(CompileError::ComputeLimitExceeded(
             String::from("X threads"),
-            x as _,
+            node_template.threads.0 as _,
             MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
         ));
     }
-    if y > MAX_COMPUTE_WORKGROUPS_PER_DIMENSION {
+    if node_template.threads.1 > MAX_COMPUTE_WORKGROUPS_PER_DIMENSION {
         return Err(CompileError::ComputeLimitExceeded(
             String::from("Y threads"),
-            y as _,
+            node_template.threads.1 as _,
             MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
         ));
     }
-    if z > MAX_COMPUTE_WORKGROUPS_PER_DIMENSION {
+    if node_template.threads.2 > MAX_COMPUTE_WORKGROUPS_PER_DIMENSION {
         return Err(CompileError::ComputeLimitExceeded(
             String::from("Z threads"),
-            z as _,
+            node_template.threads.2 as _,
             MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
         ));
     }
 
+    // Determine (default) scalar data type to use
+    context.insert("scalar_type", node_template.scalar_type.wgsl_type_name());
+    context.insert("scalar_stride", &node_template.scalar_type.stride());
+    context.insert(
+        "vec4_stride",
+        &(MultiType::Vec(node_template.scalar_type, 4).stride()),
+    );
+    context.insert(
+        "mat4x4_stride",
+        &(MultiType::Mat(node_template.scalar_type, 4, 4).stride()),
+    );
+    context.insert("mat3x3_stride", &(48));
+
+    // Render template
+    let shader = TEMPLATES
+        .render(node_template.template, &context)
+        .expect("failed to render shader");
+
     Ok(CompiledNode {
         shader,
-        threads: (x, y, z),
+        threads: node_template.threads,
     })
 }
 
diff --git a/wonnx/src/gpu.rs b/wonnx/src/gpu.rs
index 0e630093..9c662f9b 100644
--- a/wonnx/src/gpu.rs
+++ b/wonnx/src/gpu.rs
@@ -12,7 +12,7 @@ use crate::{
     ir::{Node, NodeDefinition, NodeIdentifier, OperatorDefinition},
     onnx::TensorProto,
     resource::{self, resize},
-    utils::{ceil, Shape, MINIMUM_BUFFER_SIZE_BYTES},
+    utils::{ceil, DataTypeError, InputTensor, ScalarType, Shape, MINIMUM_BUFFER_SIZE_BYTES},
 };
 
 /// The maximum number of bindings in a binding group (defined by wgpu)
@@ -66,6 +66,9 @@ pub enum GpuError {
 
     #[error("node output not found: index {0}")]
     OutputMissing(usize),
+
+    #[error("scalar type error: {0}")]
+    ScalarType(#[from] DataTypeError),
 }
 
 enum InferenceOutput {
@@ -208,9 +211,13 @@ impl GpuModel {
                     gpu_op
                 }
                 NodeDefinition::Tensor(tensor_def) => {
-                    let tensor_buffer = Arc::new(tensor_def.buffer(&self.device, outputs_readable));
+                    let tensor_buffer =
+                        Arc::new(tensor_def.buffer(&self.device, outputs_readable)?);
                     output_tensors.push(GpuTensor {
-                        shape: Shape::from(tensor_def.get_dims()),
+                        shape: Shape::from(
+                            ScalarType::from_i32(tensor_def.get_data_type())?,
+                            tensor_def.get_dims(),
+                        ),
                         buffer: tensor_buffer.clone(),
                     });
                     GpuStep::Initializer(tensor_buffer)
@@ -222,10 +229,16 @@ impl GpuModel {
                         );
                     }
 
-                    let input_shape = input_def.get_shape();
+                    let input_shape = input_def.get_shape()?;
+                    log::info!(
+                        "creating input buffer for {} shape {} size {}",
+                        input_def.get_name(),
+                        input_shape,
+                        input_shape.buffer_bytes()
+                    );
                     let input_buffer = Arc::new(resource::buffer(
                         &self.device,
-                        input_shape.buffer_len() as _,
+                        input_shape.buffer_bytes(),
                         input_def.get_name(),
                         BufferUsages::STORAGE | BufferUsages::COPY_DST,
                     ));
@@ -254,9 +267,9 @@ impl GpuModel {
     }
 
     /// Perform inference using this model and the specified inference inputs.
-    pub async fn infer(
+    pub async fn infer<'a>(
         &self,
-        inference_inputs: &HashMap<String, &[f32]>,
+        inference_inputs: &HashMap<String, InputTensor<'a>>,
     ) -> Result<HashMap<String, Vec<f32>>, GpuError> {
         log::info!("encode inference steps");
         let mut encoder = self
@@ -272,9 +285,9 @@ impl GpuModel {
     }
 
     /// Reads the relevant buffers for the requested inference outputs
-    async fn read_outputs(
+    async fn read_outputs<'a>(
         &self,
-        inference_inputs: &HashMap<String, &[f32]>,
+        inference_inputs: &HashMap<String, InputTensor<'a>>,
     ) -> Result<HashMap<String, Vec<f32>>, GpuError> {
         let mut output_data = HashMap::new();
 
@@ -283,7 +296,10 @@ impl GpuModel {
                 output_name.to_string(),
                 match output_source {
                     InferenceOutput::InferenceInput(input_name) => {
-                        inference_inputs[input_name].to_vec()
+                        match inference_inputs[input_name] {
+                            InputTensor::F32(v) => v.to_vec(),
+                            InputTensor::I32(v) => v.iter().map(|f| (*f) as f32).collect(),
+                        }
                     }
                     InferenceOutput::Tensor(tensor) => tensor.read_to_vec(&self.device).await?,
                 },
@@ -295,13 +311,13 @@ impl GpuModel {
 }
 
 trait TensorProtoExtra {
-    fn buffer(&self, device: &wgpu::Device, readable: bool) -> Buffer;
+    fn buffer(&self, device: &wgpu::Device, readable: bool) -> Result<Buffer, GpuError>;
 }
 
 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());
+    fn buffer(&self, device: &wgpu::Device, readable: bool) -> Result<Buffer, GpuError> {
+        let input_shape = Shape::from(ScalarType::from_i32(self.get_data_type())?, self.get_dims());
         log::info!(
             "creating tensor buffer {} shape {}",
             self.get_name(),
@@ -321,13 +337,13 @@ impl TensorProtoExtra for TensorProto {
         };
 
         // Do not create buffers that are too small
-        if raw_data.len() < MINIMUM_BUFFER_SIZE_BYTES as _ {
+        Ok(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)
-        }
+        })
     }
 }
 
@@ -380,7 +396,7 @@ impl<'model> OperatorDefinition<'model> {
 
                 let buffer = Arc::new(resource::buffer(
                     device,
-                    value_shape.buffer_len() as _,
+                    value_shape.buffer_bytes(),
                     output_name.as_str(),
                     buffer_usage,
                 ));
@@ -466,11 +482,11 @@ impl<'model> OperatorDefinition<'model> {
 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(
+    fn encode<'a>(
         &self,
         queue: &wgpu::Queue,
         encoder: &mut CommandEncoder,
-        inputs: &HashMap<String, &[f32]>,
+        inputs: &HashMap<String, InputTensor<'a>>,
     ) -> Result<(), GpuError> {
         match self {
             GpuStep::None | GpuStep::Forward(_) | GpuStep::Initializer(_) => {
@@ -484,11 +500,24 @@ impl GpuStep {
                     .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())),
-                );
+
+                match input_data {
+                    InputTensor::F32(float_input) => {
+                        queue.write_buffer(
+                            input_buffer,
+                            0,
+                            bytemuck::cast_slice(&resize(float_input.to_vec())),
+                        );
+                    }
+                    InputTensor::I32(int_input) => {
+                        queue.write_buffer(
+                            input_buffer,
+                            0,
+                            bytemuck::cast_slice(&resize(int_input.to_vec())),
+                        );
+                    }
+                }
+
                 Ok(())
             }
             GpuStep::Operator {
@@ -524,7 +553,19 @@ impl GpuTensor {
         // 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();
+        let result = match self.shape.data_type {
+            ScalarType::F32 => bytemuck::cast_slice(&output_data)[..output_buffer_size].to_vec(),
+            ScalarType::I32 => {
+                let result_ints: Vec<i32> =
+                    bytemuck::cast_slice(&output_data)[..output_buffer_size].to_vec();
+                result_ints.iter().map(|i| *i as f32).collect()
+            }
+            ScalarType::I64 => {
+                let result_ints: Vec<i64> =
+                    bytemuck::cast_slice(&output_data)[..output_buffer_size].to_vec();
+                result_ints.iter().map(|i| *i as f32).collect()
+            }
+        };
         drop(output_data);
         self.buffer.unmap();
         Ok(result)
diff --git a/wonnx/src/ir.rs b/wonnx/src/ir.rs
index b3d56ab2..35989006 100644
--- a/wonnx/src/ir.rs
+++ b/wonnx/src/ir.rs
@@ -1,5 +1,5 @@
 use crate::onnx::{ModelProto, NodeProto, TensorProto, ValueInfoProto};
-use crate::utils::Shape;
+use crate::utils::{DataTypeError, ScalarType, Shape};
 use std::borrow::Cow;
 use std::fmt::Debug;
 use std::hash::Hash;
@@ -65,6 +65,9 @@ pub enum IrError {
         target_node_name: String,
         input_name: String,
     },
+
+    #[error("issue with data types: {0}")]
+    Type(#[from] DataTypeError),
 }
 
 impl<'m> NodeDefinition<'m> {
@@ -180,13 +183,13 @@ impl<'model> Node<'model> {
         // 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());
+            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());
+                value_shapes.insert(output_name, vi.get_shape()?);
             }
         }
 
@@ -280,14 +283,17 @@ impl<'model> Node<'model> {
         }))
     }
 
-    pub fn output_shape(&self, output_index: usize) -> Shape {
-        match (&self.definition, output_index) {
+    pub fn output_shape(&self, output_index: usize) -> Result<Shape, IrError> {
+        Ok(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::Tensor(tensor_proto), 0) => Shape::from(
+                ScalarType::from_i32(tensor_proto.get_data_type())?,
+                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),
-        }
+        })
     }
 }
 
diff --git a/wonnx/src/lib.rs b/wonnx/src/lib.rs
index 769a9eed..af84215c 100644
--- a/wonnx/src/lib.rs
+++ b/wonnx/src/lib.rs
@@ -17,6 +17,7 @@ use protobuf::{self, Message, ProtobufError};
 use std::collections::HashMap;
 use std::path::Path;
 use std::result::Result;
+use utils::{DataTypeError, InputTensor};
 
 use crate::gpu::GpuModel;
 use thiserror::Error;
@@ -28,6 +29,12 @@ pub enum WonnxError {
 
     #[error("error executing the model: {0}")]
     SessionError(#[from] SessionError),
+
+    #[error("error in intermediate representation: {0}")]
+    IrError(#[from] IrError),
+
+    #[error("error in data types: {0}")]
+    TypeError(#[from] DataTypeError),
 }
 
 /// Creates a new session connected to the GPU.
@@ -134,9 +141,9 @@ impl Session {
     }
 
     /// Perform inference given the inputs provided and return all the outputs the model was compiled to return.
-    pub async fn run(
+    pub async fn run<'a>(
         &self,
-        inputs: &HashMap<String, &[f32]>,
+        inputs: &HashMap<String, InputTensor<'a>>,
     ) -> Result<HashMap<String, Vec<f32>>, SessionError> {
         Ok(self.gpu_model.infer(inputs).await?)
     }
diff --git a/wonnx/src/optimizer.rs b/wonnx/src/optimizer.rs
index fb4c4f97..8fbbae11 100644
--- a/wonnx/src/optimizer.rs
+++ b/wonnx/src/optimizer.rs
@@ -1,12 +1,11 @@
-use protobuf::{ProtobufEnum, RepeatedField};
+use protobuf::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},
+    utils::{attribute, get_attribute, AttributeNotFoundError, DataTypeError, ScalarType},
 };
 
 #[derive(Debug, Error)]
@@ -18,12 +17,15 @@ pub enum OptimizerError {
     Unsupported(String),
 
     #[error("invalid data type {data_type:?} for input {input} of op {op}")]
-    InvalidDataType {
-        data_type: TensorProto_DataType,
+    InvalidInputDataType {
+        data_type: ScalarType,
         input: String,
         op: String,
     },
 
+    #[error("error with data type: {0}")]
+    InvalidDataType(#[from] DataTypeError),
+
     #[error("required attribute not found: {0}")]
     AttributeNotFound(#[from] AttributeNotFoundError),
 }
@@ -319,10 +321,8 @@ impl<'model> Optimizer<'model> {
                                 // 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);
+                                    let data_type =
+                                        ScalarType::from_i32(tensor_proto.get_data_type())?;
 
                                     match (op, attr_name) {
                                         // Inputs that need to be converted to an i64 attribute
@@ -330,8 +330,7 @@ impl<'model> Optimizer<'model> {
                                         | ("Resize", "roi")
                                         | ("Resize", "sizes")
                                         | ("Reshape", "shape") => match data_type {
-                                            TensorProto_DataType::INT64
-                                            | TensorProto_DataType::UNDEFINED => {
+                                            ScalarType::I64 => {
                                                 log::info!(
                                                         "transferring input {} for op {} to i64 attribute (initializer data type: {:?})",
                                                         attr_name,
@@ -345,7 +344,7 @@ impl<'model> Optimizer<'model> {
                                                 ));
                                             }
                                             _ => {
-                                                return Err(OptimizerError::InvalidDataType {
+                                                return Err(OptimizerError::InvalidInputDataType {
                                                     data_type,
                                                     input: attr_name.to_string(),
                                                     op: op.to_string(),
@@ -354,8 +353,7 @@ impl<'model> Optimizer<'model> {
                                         },
                                         // Inputs that need to be converted to an f32 attribute
                                         ("Resize", "scales") => match data_type {
-                                            TensorProto_DataType::FLOAT
-                                            | TensorProto_DataType::UNDEFINED => {
+                                            ScalarType::F32 => {
                                                 log::info!(
                                                         "transferring input {} for op {} to f32 attribute (initializer data type: {:?})",
                                                         attr_name,
@@ -370,7 +368,7 @@ impl<'model> Optimizer<'model> {
                                                 ));
                                             }
                                             _ => {
-                                                return Err(OptimizerError::InvalidDataType {
+                                                return Err(OptimizerError::InvalidInputDataType {
                                                     data_type,
                                                     input: attr_name.to_string(),
                                                     op: op.to_string(),
diff --git a/wonnx/src/resource.rs b/wonnx/src/resource.rs
index a02000ba..f58e9f76 100644
--- a/wonnx/src/resource.rs
+++ b/wonnx/src/resource.rs
@@ -31,8 +31,13 @@ pub fn create_buffer_init<T: Clone + bytemuck::Pod>(
     })
 }
 
-pub fn buffer(device: &wgpu::Device, size: u64, name: &str, usage: BufferUsages) -> wgpu::Buffer {
-    let slice_size = usize::max(16, size as usize * std::mem::size_of::<f32>());
+pub fn buffer(
+    device: &wgpu::Device,
+    requested_size_bytes: usize,
+    name: &str,
+    usage: BufferUsages,
+) -> wgpu::Buffer {
+    let slice_size = usize::max(16, requested_size_bytes);
     let size = slice_size as wgpu::BufferAddress;
     device.create_buffer(&wgpu::BufferDescriptor {
         label: Some(name),
diff --git a/wonnx/src/utils.rs b/wonnx/src/utils.rs
index d6b34022..7ce9fbe1 100644
--- a/wonnx/src/utils.rs
+++ b/wonnx/src/utils.rs
@@ -1,8 +1,9 @@
+use protobuf::ProtobufEnum;
 use protobuf::RepeatedField;
-use serde_derive::Serialize;
 
 use crate::onnx;
 use crate::onnx::OperatorSetIdProto;
+use crate::onnx::TensorProto_DataType;
 use crate::onnx::ValueInfoProto;
 use std::convert::From;
 use std::convert::Into;
@@ -14,16 +15,17 @@ use thiserror::Error;
 * error: buffer binding size X is less than minimum 64" in Device::create_bind_group */
 pub const MINIMUM_BUFFER_SIZE_BYTES: u64 = 64;
 
-#[derive(Debug, Serialize, Clone)]
-#[serde(transparent)]
+#[derive(Debug, Clone)]
 pub struct Shape {
     pub dims: Vec<u64>,
+    pub data_type: ScalarType,
 }
 
 impl Shape {
-    pub fn from(ds: &[i64]) -> Shape {
+    pub fn from(data_type: ScalarType, dims: &[i64]) -> Shape {
         Shape {
-            dims: ds.iter().map(|x| *x as u64).collect(),
+            data_type,
+            dims: dims.iter().map(|x| *x as u64).collect(),
         }
     }
 
@@ -39,9 +41,9 @@ impl Shape {
         self.dims.iter().product()
     }
 
-    pub fn buffer_len(&self) -> u64 {
+    pub fn buffer_bytes(&self) -> usize {
         // Dividing by 4 since that is the size of an f32 in bytes
-        self.element_count().max(MINIMUM_BUFFER_SIZE_BYTES / 4)
+        (self.element_count() as usize) * self.data_type.stride()
     }
 
     pub fn dim(&self, idx: usize) -> u64 {
@@ -59,27 +61,92 @@ impl Shape {
     }
 }
 
-/** Represents a WGSL data type that can be used in a shader. The larger the data type, the more efficiently the GPU can
- * perform operations. However, large data types require the size of the data that is being worked on to be a multiple
- * of the data type (e.g. a vec4 can be used to work on a vector of 256 elements, but when used on a vector of 255 elements
- * calculation would overflow if the shader doesn't take this into account). The WGSL declaration looks like:
- *
- * struct Block {
- *    data: [[stride( dt.size_bytes() )]] dt.wgsl_type_name();
- * };
- */
-pub enum DataType {
+pub enum InputTensor<'a> {
+    F32(&'a [f32]),
+    I32(&'a [i32]),
+}
+
+#[derive(Error, Debug)]
+pub enum DataTypeError {
+    #[error("the ONNX scalar data type '{0:?}' is not supported")]
+    NotSupported(TensorProto_DataType),
+
+    #[error("the ONNX data type '{0}' is not recognized")]
+    NotRecognized(i32),
+
+    #[error("type is undefined")]
+    Undefined,
+}
+
+/// Data type for a single number
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum ScalarType {
     F32,
-    Vec(usize),
+    I64,
+    I32,
 }
 
-impl DataType {
+impl ScalarType {
+    pub fn from_i32(onnx: i32) -> Result<ScalarType, DataTypeError> {
+        let onnx_dt =
+            TensorProto_DataType::from_i32(onnx).ok_or(DataTypeError::NotRecognized(onnx))?;
+        Self::from(onnx_dt)
+    }
+
+    pub fn from(onnx: TensorProto_DataType) -> Result<ScalarType, DataTypeError> {
+        Ok(match onnx {
+            TensorProto_DataType::FLOAT => ScalarType::F32,
+            TensorProto_DataType::INT64 => ScalarType::I64,
+            TensorProto_DataType::INT32 => ScalarType::I32,
+            _ => return Err(DataTypeError::NotSupported(onnx)),
+        })
+    }
+
+    pub fn stride(&self) -> usize {
+        match self {
+            ScalarType::F32 => 4,
+            ScalarType::I32 => 4,
+            ScalarType::I64 => 8,
+        }
+    }
+
+    pub fn wgsl_type_name(&self) -> &'static str {
+        match self {
+            ScalarType::F32 => "f32",
+            ScalarType::I32 => "i32",
+            ScalarType::I64 => "i64",
+        }
+    }
+}
+
+impl Display for ScalarType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.wgsl_type_name())
+    }
+}
+
+/// Represents a WGSL data type that can be used in a shader to perform an operation on multiple scalars at once. The
+/// larger the data type, the more efficiently the GPU can perform operations. However, large data types require the size
+/// of the data that is being worked on to be a multiple of the data type (e.g. a vec4 can be used to work on a vector of
+/// 256 elements, but when used on a vector of 255 elements calculation would overflow if the shader doesn't take this
+/// into account). The WGSL declaration looks like:
+///
+/// struct Block {
+///     data: [[stride( dt.size_bytes() )]] dt.wgsl_type_name();
+/// };
+pub enum MultiType {
+    Scalar(ScalarType),
+    Vec(ScalarType, usize),
+    Mat(ScalarType, usize, usize),
+}
+
+impl MultiType {
     /// Determine the appropriate data type given the data size
-    pub fn for_size(n: usize) -> DataType {
+    pub fn for_size(n: usize, scalar: ScalarType) -> MultiType {
         let d = num::integer::gcd(n, 4);
         match d {
-            1 => DataType::F32,
-            2 | 4 => DataType::Vec(d),
+            1 => MultiType::Scalar(scalar),
+            2 | 4 => MultiType::Vec(scalar, d),
             /* 3 can't occur here because it is not a divisor of 4. Even so, we wouldn't be able to use vec3, because
             its stride is 16 instead of the expected 12, which would require padding to work properly. */
             _ => unreachable!(),
@@ -87,26 +154,33 @@ impl DataType {
     }
 
     /// Size (in bytes) of the data type (useful for setting the 'stride' in WGSL)
-    pub fn size_bytes(&self) -> usize {
+    pub fn stride(&self) -> usize {
         match self {
-            DataType::F32 => 4,
-            DataType::Vec(n) => 4 * n,
+            MultiType::Scalar(s) => s.stride(),
+
+            // FIXME: this may not always be right!
+            MultiType::Vec(st, n) => st.stride() * n,
+            MultiType::Mat(st, w, h) => st.stride() * w * h,
         }
     }
 
     /// Name of this data type in WGSL
     pub fn wgsl_type_name(&self) -> String {
         match self {
-            DataType::F32 => "f32".to_string(),
-            DataType::Vec(n) => format!("vec{}<f32>", n),
+            MultiType::Scalar(s) => s.wgsl_type_name().to_string(),
+            MultiType::Vec(st, n) => format!("vec{}<{}>", n, st.wgsl_type_name()),
+            MultiType::Mat(st, w, h) => format!("mat{}x{}<{}>", w, h, st.wgsl_type_name()),
         }
     }
 
     /// The number of elements in this data type
     pub fn elements(&self) -> usize {
         match self {
-            DataType::F32 => 1,
-            DataType::Vec(n) => *n,
+            MultiType::Scalar(_) => 1,
+
+            // FIXME: this may not always be right
+            MultiType::Vec(_, n) => *n,
+            &MultiType::Mat(_, w, h) => w * h,
         }
     }
 }
@@ -158,22 +232,41 @@ pub fn ceil(num: u64, div: u64) -> u64 {
 }
 
 impl ValueInfoProto {
-    pub fn get_shape(&self) -> Shape {
-        Shape::from(
-            self.get_field_type()
-                .get_tensor_type()
-                .get_shape()
-                .get_dim()
-                .iter()
-                .map(|x| x.get_dim_value() as i64)
-                .collect::<Vec<i64>>()
-                .as_slice(),
-        )
+    pub fn get_shape(&self) -> Result<Shape, DataTypeError> {
+        Ok(match &self.get_field_type().value {
+            Some(t) => match t {
+                onnx::TypeProto_oneof_value::tensor_type(tensor_proto) => Shape::from(
+                    ScalarType::from_i32(tensor_proto.get_elem_type())?,
+                    self.get_field_type()
+                        .get_tensor_type()
+                        .get_shape()
+                        .get_dim()
+                        .iter()
+                        .map(|x| x.get_dim_value() as i64)
+                        .collect::<Vec<i64>>()
+                        .as_slice(),
+                ),
+                onnx::TypeProto_oneof_value::sequence_type(_) => todo!(),
+                onnx::TypeProto_oneof_value::map_type(_) => todo!(),
+                onnx::TypeProto_oneof_value::optional_type(_) => todo!(),
+                onnx::TypeProto_oneof_value::sparse_tensor_type(_) => todo!(),
+            },
+            None => return Err(DataTypeError::Undefined),
+        })
     }
 }
 
-// TODO: Make dimension optional
+/// Shorthand method to define an ONNX tensor with the specified name and shape (data type is f32)
 pub fn tensor(name: &str, dimensions: &[i64]) -> onnx::ValueInfoProto {
+    tensor_of_type(name, dimensions, TensorProto_DataType::FLOAT)
+}
+
+/// Shorthand method to define an ONNX tensor with the specified name, shape and data type
+pub fn tensor_of_type(
+    name: &str,
+    dimensions: &[i64],
+    data_type: TensorProto_DataType,
+) -> onnx::ValueInfoProto {
     let mut dim_value = vec![];
     for dimension in dimensions {
         let mut dim_channel = onnx::TensorShapeProto_Dimension::new();
@@ -185,7 +278,7 @@ pub fn tensor(name: &str, dimensions: &[i64]) -> onnx::ValueInfoProto {
     shape_tensor_proto.set_dim(protobuf::RepeatedField::from(dim_value));
 
     let mut type_proto_tensor = onnx::TypeProto_Tensor::new();
-    type_proto_tensor.set_elem_type(1);
+    type_proto_tensor.set_elem_type(data_type.value());
     type_proto_tensor.set_shape(shape_tensor_proto);
 
     let mut type_proto = onnx::TypeProto::new();
@@ -202,6 +295,7 @@ pub fn tensor(name: &str, dimensions: &[i64]) -> onnx::ValueInfoProto {
 pub fn initializer(name: &str, data: Vec<f32>) -> onnx::TensorProto {
     let mut initializer = crate::onnx::TensorProto::new();
     initializer.set_name(name.to_string());
+    initializer.set_data_type(1); // FLOAT
     initializer.set_float_data(data);
     initializer
 }
@@ -345,7 +439,7 @@ impl From<onnx::AttributeProto> for String {
 
 #[cfg(test)]
 mod tests {
-    use crate::utils::{attribute, graph, initializer, model, node, tensor};
+    use crate::utils::{attribute, graph, initializer, model, node, tensor, InputTensor};
 
     #[test]
     fn test_model() {
@@ -356,7 +450,7 @@ mod tests {
         let mut input_data = std::collections::HashMap::new();
 
         let data: Vec<f32> = (0..25).map(|x| x as f32).collect();
-        input_data.insert("X".to_string(), data.as_slice());
+        input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
         // ONNX INPUTS
         let shape = vec![1, c, n, n];
diff --git a/wonnx/templates/endomorphism/activation.wgsl b/wonnx/templates/endomorphism/activation.wgsl
index e8247c26..d3f5eb67 100644
--- a/wonnx/templates/endomorphism/activation.wgsl
+++ b/wonnx/templates/endomorphism/activation.wgsl
@@ -8,11 +8,10 @@ var<storage, write> output_0: ArrayVector;
 
 [[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 = "input_0.data[gidx]" %}
-{% set activation_output = "output_0.data[gidx]" %}
-{% set activation_type = op_type %}
-{%- include "snippets/activation_vec.wgsl" -%}
+	let gidx = global_id.x;
 
+	{% set activation_input = "input_0.data[gidx]" %}
+	{% set activation_output = "output_0.data[gidx]" %}
+	{% set activation_type = op_type %}
+	{%- include "snippets/activation_vec.wgsl" -%}
 }
\ No newline at end of file
diff --git a/wonnx/templates/endomorphism/arithmetic.wgsl b/wonnx/templates/endomorphism/arithmetic.wgsl
index d90159f5..babbd698 100644
--- a/wonnx/templates/endomorphism/arithmetic.wgsl
+++ b/wonnx/templates/endomorphism/arithmetic.wgsl
@@ -20,10 +20,18 @@ var<storage, write> output_0: ArrayVector;
 
 [[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 i_lens | length == 2 %}
-    output_0.data[gidx] = input_0.data[gidx] {{ op_type }} input_1.data[gidx];
-{% else %}
-    output_0.data[gidx] = input_0.data[gidx] {{ op_type }} vec4<f32>({{ coefficient }}., {{ coefficient }}., {{ coefficient }}., {{ coefficient }}.);
-{% endif %}
+	let gidx = global_id.x;
+
+	{% if i_lens | length == 2 %}
+		output_0.data[gidx] = input_0.data[gidx] {{ op_type }} input_1.data[gidx];
+
+	{% else %}
+		output_0.data[gidx] = input_0.data[gidx] {{ op_type }} Vec4(
+			Scalar({{ coefficient }}), 
+			Scalar({{ coefficient }}),
+			Scalar({{ coefficient }}),
+			Scalar({{ coefficient }})
+		);
+		
+	{% endif %}
 }
\ No newline at end of file
diff --git a/wonnx/templates/endomorphism/batchnormalization.wgsl b/wonnx/templates/endomorphism/batchnormalization.wgsl
index c3b3c299..e483e64a 100644
--- a/wonnx/templates/endomorphism/batchnormalization.wgsl
+++ b/wonnx/templates/endomorphism/batchnormalization.wgsl
@@ -2,8 +2,8 @@
 {%- include "structs.wgsl" -%}
 
 struct Block {
-    data: [[stride({{ elem_stride }})]] array<{{ elem_type }}>;
-}; 
+	data: [[stride({{ elem_stride }})]] array<{{ elem_type }}>;
+};
 
 // X (input)
 [[group(0), binding(0)]]
@@ -31,16 +31,16 @@ var<storage, write> output_0: Block;
 
 [[stage(compute), workgroup_size(1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-    let channel = global_id.y;
-    let batch = global_id.z;
-    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 = 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];
-
-    output_0.data[index] = (x - channel_mean) / sqrt(channel_var + {{ epsilon }}) * channel_scale + channel_bias;
+	let channel = global_id.y;
+	let batch = global_id.z;
+	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 = 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];
+
+	output_0.data[index] = (x - channel_mean) / sqrt(channel_var + {{ epsilon }}) * channel_scale + channel_bias;
 }
\ No newline at end of file
diff --git a/wonnx/templates/endomorphism/cast.wgsl b/wonnx/templates/endomorphism/cast.wgsl
new file mode 100644
index 00000000..5fd02320
--- /dev/null
+++ b/wonnx/templates/endomorphism/cast.wgsl
@@ -0,0 +1,16 @@
+{%- include "structs.wgsl" -%}
+[[group(0), binding(0)]]
+var<storage, read> input_0: ArrayVector;
+
+struct OutputArrayVector {
+	data: [[stride({{ vec4_stride }})]] array<vec4<{{ cast_to_type }}>>;
+}; 
+
+[[group(0), binding(1)]]
+var<storage, write> output_0: OutputArrayVector;
+
+[[stage(compute), workgroup_size({{ workgroup_size_x }})]]
+fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
+    let gidx = global_id.x;
+    output_0.data[gidx] = vec4<{{ cast_to_type }}>(input_0.data[gidx]);
+}
\ No newline at end of file
diff --git a/wonnx/templates/endomorphism/map.wgsl b/wonnx/templates/endomorphism/map.wgsl
index c2f12525..af1dc266 100644
--- a/wonnx/templates/endomorphism/map.wgsl
+++ b/wonnx/templates/endomorphism/map.wgsl
@@ -7,6 +7,14 @@ var<storage, write> output_0: ArrayVector;
 
 [[stage(compute), workgroup_size({{ workgroup_size_x }})]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-    let gidx = global_id.x;
-    output_0.data[gidx] = {{ op_type | lower }}(input_0.data[gidx]);
+	let gidx = global_id.x;
+	
+	{% if op_type == "Reciprocal" %}
+		let one = Vec4(Scalar(1), Scalar(1), Scalar(1), Scalar(1));
+		output_0.data[gidx] = one / (input_0.data[gidx]);
+
+	{% else %}
+		output_0.data[gidx] = {{ op_type | lower }}(input_0.data[gidx]);
+
+	{% endif %}
 }
\ No newline at end of file
diff --git a/wonnx/templates/endomorphism/softmax.wgsl b/wonnx/templates/endomorphism/softmax.wgsl
index 01fe8fbc..4cb7ce58 100644
--- a/wonnx/templates/endomorphism/softmax.wgsl
+++ b/wonnx/templates/endomorphism/softmax.wgsl
@@ -1,4 +1,3 @@
-
 {%- include "structs.wgsl" -%}
 
 [[group(0), binding(0)]]
@@ -12,28 +11,28 @@ fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 	let gidx = global_id.x;
 
 	{% if op_type == "Softmax" %}
-	// Softmax = exp(input - max(input)) / sum(exp(input - max(input)))
-	
-	// First, determine max(input)
-	// WGSL doesn't have a way to write -Infinity (https://github.com/gpuweb/gpuweb/issues/1769)
-	// 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 = input_0.data[gidx + k];
-		max_element = max(max_element, element);
-	}
+		// Softmax = exp(input - max(input)) / sum(exp(input - max(input)))
+		
+		// First, determine max(input)
+		// WGSL doesn't have a way to write -Infinity (https://github.com/gpuweb/gpuweb/issues/1769)
+		// Therefore we use log(0) instead which returns -Infinity
+		var max_element: Scalar = log(Scalar(0));
+		for(var k: u32 = 0u; k < {{ i_lens[0] }}u; k = k + 1u) {
+			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 = input_0.data[gidx + k];
-		sum  = sum + exp(element - max_element);
-	}
+		// Calculate sum(exp(input - max(input)))
+		var sum: Scalar = Scalar(0);
+		for(var k: u32 = 0u; k < {{ i_lens[0] }}u; k = k + 1u) {
+			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 = input_0.data[gidx + k];
-		output_0.data[gidx + k] = exp(element - max_element) / sum;
-	}
+		// Calculate elements and write to output
+		for(var k: u32 = 0u; k < {{ i_lens[0] }}u; k = k + 1u) {
+			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/wonnx/templates/matrix/concat.wgsl b/wonnx/templates/matrix/concat.wgsl
index 7fa50e87..803bc248 100644
--- a/wonnx/templates/matrix/concat.wgsl
+++ b/wonnx/templates/matrix/concat.wgsl
@@ -1,29 +1,32 @@
 
 {%- include "structs.wgsl" -%}
 
-
 {% for input in i_lens %}
+
 [[group({{ loop.index0 / 4 | int }}), binding({{ loop.index0 % 4}})]]
 var<storage, read> input_{{ loop.index0 }}: Array;
+
 {% endfor %}
 
 {% set binding_len = i_lens | length %}
 [[group({{ binding_len  / 4 | int }}), binding({{ binding_len % 4 }})]]
 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 i_lens %}
-{% if loop.first %}
-    if (gidx < {{ i_lens[0] }}u) {
-	    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)) {
-	    output_0.data[gidx] = input_{{ loop.index0 }}.data[gidx - {{ cum_len | nth(n=loop.index0 -1) }}u];
-    }
-{% endif %}
-{% endfor %}
+	let gidx = global_id.x;
+	
+	{% for input in i_lens %}
+		{% if loop.first %}
+			if (gidx < {{ i_lens[0] }}u) {
+				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)) {
+				output_0.data[gidx] = input_{{ loop.index0 }}.data[gidx - {{ cum_len | nth(n=loop.index0 -1) }}u];
+			}
+			
+		{% endif %}
+	{% endfor %}
 }
diff --git a/wonnx/templates/matrix/gemm.wgsl b/wonnx/templates/matrix/gemm.wgsl
index 211a4059..27901360 100644
--- a/wonnx/templates/matrix/gemm.wgsl
+++ b/wonnx/templates/matrix/gemm.wgsl
@@ -19,47 +19,53 @@ var<storage, write> output_0: ArrayVector;
 
 [[stage(compute), workgroup_size(1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-    let y = global_id.x % {{ i_shape[1][1] / 4 | int }}u;
-    let x = global_id.x / {{ i_shape[1][1] / 4 | int }}u;
-    let index = x * {{ i_shape[1][1] }}u + y;
+	let y = global_id.x % {{ i_shape[1][1] / 4 | int }}u;
+	let x = global_id.x / {{ i_shape[1][1] / 4 | int }}u;
+	let index = x * {{ i_shape[1][1] }}u + y;
 
-    var tmpsum = mat4x4<f32>(vec4<f32>(0.0, 0.0, 0.0, 0.0), vec4<f32>(0.0, 0.0, 0.0, 0.0), vec4<f32>(0.0, 0.0, 0.0, 0.0), vec4<f32>(0.0, 0.0, 0.0, 0.0));
-    var product = mat4x4<f32>(vec4<f32>(0.0, 0.0, 0.0, 0.0), vec4<f32>(0.0, 0.0, 0.0, 0.0), vec4<f32>(0.0, 0.0, 0.0, 0.0), vec4<f32>(0.0, 0.0, 0.0, 0.0));
+	let zero = Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0));
+	var tmpsum = Mat4x4(zero, zero, zero, zero);
+	var product = Mat4x4(zero, zero, zero, zero);
 
-    for(var k: u32 = 0u; k < {{ i_shape[0][1] / 4 | int }}u; k = k + 1u) {
-        let index_left = x * {{ i_shape[0][1] }}u + k; 
-        let index_right = k * {{ i_shape[0][1] }}u + y; 
+	for(var k: u32 = 0u; k < {{ i_shape[0][1] / 4 | int }}u; k = k + 1u) {
+		let index_left = x * {{ i_shape[0][1] }}u + k; 
+		let index_right = k * {{ i_shape[0][1] }}u + y; 
 
-        let mat_left = mat4x4<f32>(
-                              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>(
-                              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],
-                          );
+		let mat_left = Mat4x4(
+			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(
+			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;
+		product = mat_right * mat_left;
 	
-        for(var index_mat: u32 = 0u; index_mat < 4u; index_mat = index_mat + 1u) {
-	        tmpsum[index_mat] = tmpsum[index_mat] + product[index_mat];
+		for(var index_mat: u32 = 0u; index_mat < 4u; index_mat = index_mat + 1u) {
+			tmpsum[index_mat] = tmpsum[index_mat] + product[index_mat];
+		}
 	}
-    }
-    
-{%- 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) {
-            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) {
-            output_0.data[index + index_mat * {{ i_shape[1][1] / 4 | int }}u] = {%- if alpha != 1 -%} {{ alpha | float }} * {%- endif -%}tmpsum[index_mat];
-    }       
-{%- endif -%}  
+	
+	{%- if i_lens | length == 3 -%}
+		let bias_row = input_2.data[x]; 
+		var bias = transpose(Mat4x4(bias_row, bias_row, bias_row, bias_row));
+		for(var index_mat: u32 = 0u; index_mat < 4u; index_mat = index_mat + 1u) {
+			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) {
+			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/wonnx/templates/matrix/gemm_1.wgsl b/wonnx/templates/matrix/gemm_1.wgsl
index 7eb91e0f..8b2b1521 100644
--- a/wonnx/templates/matrix/gemm_1.wgsl
+++ b/wonnx/templates/matrix/gemm_1.wgsl
@@ -21,32 +21,34 @@ var<storage, write> output_0: Array;
 
 [[stage(compute), workgroup_size(1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-    let gidx = global_id.x;
+	let gidx = global_id.x;
 
-    var tmpsum = 0.0;
-    var product = 0.0;
+	var tmpsum = Scalar(0);
+	var product = Scalar(0);
 
-    for(var k: u32 = 0u; k < {{ i_shape[0][1] / 4 | int }}u; k = k + 1u) {
+	for(var k: u32 = 0u; k < {{ i_shape[0][1] / 4 | int }}u; k = k + 1u) {
+		let index_left = k; 
+		let index_right = k * {{ i_shape[1][1] * 4 }}u + gidx; 
 
-        let index_left = k; 
-        let index_right = k * {{ i_shape[1][1] * 4 }}u + gidx; 
+		let vec_left = input_0.data[index_left];
 
-        let vec_left = input_0.data[index_left];
-
-        let vec_right = vec4<f32>(
-                              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],
-                          );
+		let vec_right = Vec4(
+			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);
+		product = dot(vec_left, vec_right);
+		tmpsum = tmpsum + product;
+	}
 	
-	    tmpsum = tmpsum + product;
-    }
-    
-    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 -%};
+	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/wonnx/templates/matrix/resize.wgsl b/wonnx/templates/matrix/resize.wgsl
index 798fce2b..0c3afb93 100644
--- a/wonnx/templates/matrix/resize.wgsl
+++ b/wonnx/templates/matrix/resize.wgsl
@@ -4,38 +4,36 @@
 [[group(0), binding(0)]]
 var<storage, read> input_0: Array;
 
-
 [[group(0), binding(1)]]
 var<storage, write> output_0: Array;
 
-// resize.wgsl
 [[stage(compute), workgroup_size(256, 1, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-    let gidx = global_id.x;
-
-    if (gidx < {{ o_lens[0] }}u) {
-
-        var rest = gidx;
-        {%- for chunk in o_chunks[0] -%}
-        {% if loop.last %}
-        let d_{{ loop.index0 }} = rest; 
-        {% else %}
-        let d_{{ loop.index0 }} = rest / {{ chunk }}u; 
-        rest = gidx % {{ chunk }}u; 
-        {% endif %}
-        {%- endfor %}
-
-
-        let index = {%- for chunks in i_chunks[0] -%}
+	let gidx = global_id.x;
+
+	if (gidx < {{ o_lens[0] }}u) {
+		var rest = gidx;
+		{%- for chunk in o_chunks[0] -%}
+			{% if loop.last %}
+				let d_{{ loop.index0 }} = rest; 
+			{% else %}
+				let d_{{ loop.index0 }} = rest / {{ chunk }}u; 
+				rest = gidx % {{ chunk }}u; 
+			{% endif %}
+		{%- endfor %}
+
+		let index = 
+			{%- for chunks in i_chunks[0] -%}
 				{% set scale = scales | nth(n=loop.index0) %}
-        {%- if not loop.first %}
-				+ 
+				{%- if not loop.first %}
+					+ 
 				{%- endif -%}
-        u32(floor(
-				  (f32(d_{{ loop.index0 }}) + 0.5) / {{ scale }} - 0.5 
-			  )) * {{ chunks  }}u 
-        {%- endfor -%};
-
-	    output_0.data[gidx] = input_0.data[index];
-    }
+				u32(floor(
+					(Scalar(d_{{ loop.index0 }}) + Scalar(0.5)) / {{ scale }} - Scalar(0.5) 
+				)) * {{ chunks  }}u 
+			{%- endfor -%}
+		;
+
+		output_0.data[gidx] = input_0.data[index];
+	}
 }
diff --git a/wonnx/templates/matrix/split.wgsl b/wonnx/templates/matrix/split.wgsl
index a40c8257..f6d645ab 100644
--- a/wonnx/templates/matrix/split.wgsl
+++ b/wonnx/templates/matrix/split.wgsl
@@ -1,64 +1,63 @@
 
 {%- include "structs.wgsl" -%}
 
-
 [[group(0), binding(0)]]
 var<storage, read> input_0: Array;
 
-
 {% for output in o_lens %}
-[[group({{ loop.index / 4 | int }}), binding({{ loop.index % 4}})]]
-var<storage, write> output_{{ loop.index0 }}: Array;
+	[[group({{ loop.index / 4 | int }}), binding({{ loop.index % 4}})]]
+	var<storage, write> output_{{ loop.index0 }}: Array;
 {% endfor %}
 
-// split.wgsl
 [[stage(compute), workgroup_size(256, 1, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-    let gidx = global_id.x;
-
-    if (gidx < {{ i_lens[0] }}u) {
-
-        var rest = gidx;
-        {%- for chunks in i_chunks[0] -%}
-        {% if loop.last %}
-        let d_{{ loop.index0 }} = rest; 
-        {% else %}
-        let d_{{ loop.index0 }} = rest / {{ chunks }}u; 
-        rest = gidx % {{ chunks }}u; 
-        {% endif %}
-        {%- endfor -%}
-
-{% for output in o_lens %}
-{%- if loop.first %}
-    if (d_{{ axis }} < {{ split | first }}u) {
-
-        let index = {%- for chunk in o_chunks | first -%}
-        {%- if not loop.first %}
-        +
-        {%- endif -%}
-        d_{{ loop.index0 }} * {{ chunk }}u
-        {%- endfor -%};
-
-	    output_{{ loop.index0 }}.data[index] = input_0.data[gidx];
-    }
-{%- else %}
-{% set split_output = split | nth(n=loop.index0 -1) %}
-	if ((d_{{ axis }} >= {{ split_output }}u) && (d_{{ axis }} < {{ split | nth(n=loop.index0)}}u)) {
-
-        let index = {%- for chunk in o_chunks | nth(n=loop.index0) -%}
-        {%- if not loop.first %}
-        +
-        {%- endif -%}
-        {%- if loop.index0 == axis %}
-        (d_{{ loop.index0 }} - {{ split_output }}u) * {{ chunk }}u
-        {% else %}
-        d_{{ loop.index0 }} * {{ chunk }}u
-        {%- endif -%}
-        {%- endfor -%};
-
-	    output_{{ loop.index0 }}.data[index] = input_0.data[gidx];
-    }
-{% endif %}
-{% endfor %}
-    }
+	let gidx = global_id.x;
+
+	if (gidx < {{ i_lens[0] }}u) {
+		var rest = gidx;
+
+		{%- for chunks in i_chunks[0] -%}
+			{% if loop.last %}
+				let d_{{ loop.index0 }} = rest; 
+			{% else %}
+				let d_{{ loop.index0 }} = rest / {{ chunks }}u; 
+				rest = gidx % {{ chunks }}u; 
+			{% endif %}
+		{%- endfor -%}
+
+		{% for output in o_lens %}
+			{%- if loop.first %}
+				if (d_{{ axis }} < {{ split | first }}u) {
+					let index = 
+						{%- for chunk in o_chunks | first -%}
+							{%- if not loop.first %}
+								+
+							{%- endif -%}
+							d_{{ loop.index0 }} * {{ chunk }}u
+						{%- endfor -%}
+					;
+
+					output_{{ loop.index0 }}.data[index] = input_0.data[gidx];
+				}
+			{%- else %}
+				{% set split_output = split | nth(n=loop.index0 -1) %}
+				if ((d_{{ axis }} >= {{ split_output }}u) && (d_{{ axis }} < {{ split | nth(n=loop.index0)}}u)) {
+					let index = 
+						{%- for chunk in o_chunks | nth(n=loop.index0) -%}
+							{%- if not loop.first %}
+								+
+							{%- endif -%}
+							
+							{%- if loop.index0 == axis %}
+								(d_{{ loop.index0 }} - {{ split_output }}u) * {{ chunk }}u
+							{% else %}
+								d_{{ loop.index0 }} * {{ chunk }}u
+							{%- endif -%}
+						{%- endfor -%}
+					;
+					output_{{ loop.index0 }}.data[index] = input_0.data[gidx];
+				}
+			{% endif %}
+		{% endfor %}
+	}
 }
diff --git a/wonnx/templates/matrix/transpose.wgsl b/wonnx/templates/matrix/transpose.wgsl
index 8725976d..98447db3 100644
--- a/wonnx/templates/matrix/transpose.wgsl
+++ b/wonnx/templates/matrix/transpose.wgsl
@@ -6,30 +6,32 @@ var<storage, read> input_0: Array;
 [[group(0), binding(1)]]
 var<storage, write> output_0: Array;
 
-
 [[stage(compute), workgroup_size(256, 1, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-    let gidx = global_id.x;
-
-    if (gidx < {{ i_lens[0] }}u) {
-
-        var rest = gidx;
-        {%- for chunks in i_chunks[0] -%}
-        {% if loop.last %}
-        let d_{{ loop.index0 }} = rest; 
-        {% else %}
-        let d_{{ loop.index0 }} = rest / {{ chunks }}u; 
-        rest = gidx % {{ chunks }}u; 
-        {% endif %}
-        {%- endfor -%}
-
-        let index = {%- for perm in permuted_chunks -%}
-        {% if not loop.first %}
-        +
-        {% endif %}
-        d_{{ loop.index0 }} * {{ perm }}u
-        {%- endfor %};
-
-        output_0.data[index] = input_0.data[gidx];
-    }
+	let gidx = global_id.x;
+
+	if (gidx < {{ i_lens[0] }}u) {
+		var rest = gidx;
+
+		{%- for chunks in i_chunks[0] -%}
+			{% if loop.last %}
+				let d_{{ loop.index0 }} = rest; 
+			{% else %}
+				let d_{{ loop.index0 }} = rest / {{ chunks }}u; 
+				rest = gidx % {{ chunks }}u; 
+			{% endif %}
+		{%- endfor -%}
+
+		let index = 
+			{%- for perm in permuted_chunks -%}
+				{% if not loop.first %}
+					+
+				{% endif %}
+				
+				d_{{ loop.index0 }} * {{ perm }}u
+			{%- endfor %}
+		;
+
+		output_0.data[index] = input_0.data[gidx];
+	}
 }
\ No newline at end of file
diff --git a/wonnx/templates/pool/aggregate.wgsl b/wonnx/templates/pool/aggregate.wgsl
index cf4e1e60..06abe3d3 100644
--- a/wonnx/templates/pool/aggregate.wgsl
+++ b/wonnx/templates/pool/aggregate.wgsl
@@ -6,60 +6,55 @@ var<storage, read> input_0: Array;
 [[group(0), binding(1)]]
 var<storage, write> output_0: Array;
 
-// aggregate.wgsl
 [[stage(compute), workgroup_size(256, 1, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 	let gidx = global_id.x;
-        if (gidx < {{ o_lens[0] / 4 }}u) {
 
-	let batch = gidx / {{ o_chunks[0][0] / 4 }}u; 
-	let rest = gidx % {{ o_chunks[0][0] / 4 }}u; 
+	if (gidx < {{ o_lens[0] / 4 }}u) {
+		let batch = gidx / {{ o_chunks[0][0] / 4 }}u; 
+		let rest = gidx % {{ o_chunks[0][0] / 4 }}u; 
 
-        let m = rest / {{ o_chunks[0][1] }}u;
-        let rest = rest % {{ o_chunks[0][1] }}u;
- 
-        let y = rest / {{ o_chunks[0][2] }}u;
-        let x = rest % {{ o_chunks[0][2] }}u;
-        
-        var result = vec4<f32>(0.0, 0.0, 0.0, 0.0);
-        
-        let base_index = batch * {{ i_chunks[0][0] }}u + m * {{ i_chunks[0][1] * 4 }}u + y * {{ stride[0] }}u * {{ original_width }}u+ x * {{ stride[1] }}u;
-        for(var i: u32 = 0u; i < {{ kernel_shape[0] }}u; i = i + 1u) {
-                
-		let tmp_y = i * {{ dilation[0] }}u; 
-        
-	        for(var j: u32 = 0u; j < {{ kernel_shape[1] }}u; j = j + 1u) { 
-
-                        let tmp_x = j * {{ dilation[1] }}u;
-
-                        let tmp_index  = base_index + tmp_y * {{ original_width }}u + tmp_x;
-                        let vector = vec4<f32>(
-                                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" -%}
-
-			result = max(result, vector);
-
-{%- elif op_type == "AveragePool" -%}
-			result = result + vector;
-{%- endif -%}
-                        }
-  	        }
+		let m = rest / {{ o_chunks[0][1] }}u;
+		let rest = rest % {{ o_chunks[0][1] }}u;
+	
+		let y = rest / {{ o_chunks[0][2] }}u;
+		let x = rest % {{ o_chunks[0][2] }}u;
 		
-            
-
-{% if op_type == "AveragePool" -%}
-        result = result / {{ kernel_len }}.;
-{%- endif %}
-
-        let base_index = batch * {{ o_chunks[0][0] }}u + m * {{ o_chunks[0][1] * 4 }}u + y * {{ width }}u + x;
-        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;
+		var result = Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0));
+		
+		let base_index = batch * {{ i_chunks[0][0] }}u + m * {{ i_chunks[0][1] * 4 }}u + y * {{ stride[0] }}u * {{ original_width }}u+ x * {{ stride[1] }}u;
 
-                output_0.data[index] = result[index_vec];
-        }
-        }
+		for(var i: u32 = 0u; i < {{ kernel_shape[0] }}u; i = i + 1u) {
+			let tmp_y = i * {{ dilation[0] }}u; 
+		
+			for(var j: u32 = 0u; j < {{ kernel_shape[1] }}u; j = j + 1u) { 
+				let tmp_x = j * {{ dilation[1] }}u;
+
+				let tmp_index  = base_index + tmp_y * {{ original_width }}u + tmp_x;
+				let vector = Vec4(
+					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" -%}
+					result = max(result, vector);
+				{%- elif op_type == "AveragePool" -%}
+					result = result + vector;
+				{%- endif -%}
+			}
+		}
+
+		{% if op_type == "AveragePool" -%}
+			result = result / {{ kernel_len }}.;
+		{%- endif %}
+
+		let base_index = batch * {{ o_chunks[0][0] }}u + m * {{ o_chunks[0][1] * 4 }}u + y * {{ width }}u + x;
+
+		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;
+			output_0.data[index] = result[index_vec];
+		}
+	}
 }
diff --git a/wonnx/templates/pool/conv.wgsl b/wonnx/templates/pool/conv.wgsl
index 82c40883..3fdd2b8b 100644
--- a/wonnx/templates/pool/conv.wgsl
+++ b/wonnx/templates/pool/conv.wgsl
@@ -7,73 +7,64 @@ var<storage, read> input_0: Array;
 var<storage, read> input_1: Array;
 
 {%- if i_lens | length == 3 -%} // Bias
-[[group(0), binding(2)]]
-var<storage, read> input_2: Array;
+	[[group(0), binding(2)]]
+	var<storage, read> input_2: Array;
 
-[[group(0), binding(3)]]
-var<storage, write> output_0: Array;
+	[[group(0), binding(3)]]
+	var<storage, write> output_0: Array;
 
 {%- else -%}
-[[group(0), binding(2)]]
-var<storage, write> output_0: Array;
-{%- endif %}  
+	[[group(0), binding(2)]]
+	var<storage, write> output_0: Array;
+
+{%- endif %}
 
-// Conv.wgsl
 [[stage(compute), workgroup_size(256, 1, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 	let gidx = global_id.x;
-        if (gidx < {{ o_lens[0] }}u) {
-	let batch = gidx / {{ o_chunks[0][0] }}u; 
-	let rest = gidx % {{ o_chunks[0][0] }}u; 
-
-        let m = rest / {{ o_chunks[0][1] }}u;
-        let rest = rest % {{ o_chunks[0][1] }}u;
-        
-        let y = rest / {{ o_chunks[0][2] }}u;
-        let x = rest % {{ o_chunks[0][2] }}u;
-        
-        var result: f32 = 0.0;
-
-        let root_index = batch * {{ i_chunks[0][0] }}u;
-
-        let root_kernel_index = m * {{ kernel_channel_len }}u;
-
-        for(var c: u32 = 0u; c < {{ channel }}u; c = c + 1u) {
-            
-            let base_index = root_index + c * {{ i_chunks[0][1] }}u;
-            let base_kernel_index = root_kernel_index + c * {{ kernel_len }}u;
-
-            for(var i: u32 = 0u; i < {{ kernel_shape[0] }}u; i = i + 1u) {
-                
-		let tmp_y = y * {{ stride[0] }}u + i * {{ dilation[0] }}u - {{ pad[0] }}u; 
-
-        	if ((tmp_y < {{ original_height }}u) && (tmp_y >= 0u)) {
-        
-	        for(var j: u32 = 0u; j < {{ kernel_shape[1] }}u; j = j + 1u) { 
-
-                        let tmp_x = x * {{ stride[1] }}u + j * {{ dilation[1] }}u - {{ pad[1] }}u;
-
-                        if ((tmp_x < {{ original_width }}u) && (tmp_x >= 0u)) {
-
-                                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 = input_0.data[tmp_index] * input_1.data[index_kernel] + result;
-				
-                        }
-  	        }
+	if (gidx < {{ o_lens[0] }}u) {
+		let batch = gidx / {{ o_chunks[0][0] }}u; 
+		let rest = gidx % {{ o_chunks[0][0] }}u; 
+
+		let m = rest / {{ o_chunks[0][1] }}u;
+		let rest = rest % {{ o_chunks[0][1] }}u;
+		
+		let y = rest / {{ o_chunks[0][2] }}u;
+		let x = rest % {{ o_chunks[0][2] }}u;
+		
+		var result: Scalar = Scalar(0);
+
+		let root_index = batch * {{ i_chunks[0][0] }}u;
+		let root_kernel_index = m * {{ kernel_channel_len }}u;
+
+		for(var c: u32 = 0u; c < {{ channel }}u; c = c + 1u) {
+			let base_index = root_index + c * {{ i_chunks[0][1] }}u;
+			let base_kernel_index = root_kernel_index + c * {{ kernel_len }}u;
+
+			for(var i: u32 = 0u; i < {{ kernel_shape[0] }}u; i = i + 1u) {
+				let tmp_y = y * {{ stride[0] }}u + i * {{ dilation[0] }}u - {{ pad[0] }}u; 
+
+				if ((tmp_y < {{ original_height }}u) && (tmp_y >= 0u)) {
+					for(var j: u32 = 0u; j < {{ kernel_shape[1] }}u; j = j + 1u) { 
+						let tmp_x = x * {{ stride[1] }}u + j * {{ dilation[1] }}u - {{ pad[1] }}u;
+
+						if ((tmp_x < {{ original_width }}u) && (tmp_x >= 0u)) {
+							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 = input_0.data[tmp_index] * input_1.data[index_kernel] + result;			
+						}
+					}
+				}
+			}
 		}
-            }
-	}
-
-        {%- if i_lens | length == 3 -%}
-        result = result + input_2.data[m];
-        {%- endif -%}
 
-{% set activation_input = "result" %}
-{% set activation_output = "output_0.data[gidx]" %}
-{% set activation_type = op_type | replace(from="Conv", to="") %}
-{%- include "snippets/activation_scalar.wgsl" -%}
+		{%- if i_lens | length == 3 -%}
+			result = result + input_2.data[m];
+		{%- endif -%}
 
-    }
+		{% set activation_input = "result" %}
+		{% set activation_output = "output_0.data[gidx]" %}
+		{% set activation_type = op_type | replace(from="Conv", to="") %}
+		{%- include "snippets/activation_scalar.wgsl" -%}
+	}
 }
diff --git a/wonnx/templates/pool/conv_kernel_1.wgsl b/wonnx/templates/pool/conv_kernel_1.wgsl
index 5419916b..63ad6639 100644
--- a/wonnx/templates/pool/conv_kernel_1.wgsl
+++ b/wonnx/templates/pool/conv_kernel_1.wgsl
@@ -7,75 +7,75 @@ var<storage, read> input_0: Array;
 var<storage, read> input_1: ArrayMatrix;
 
 {%- if i_lens | length == 3 -%} // Bias
-[[group(0), binding(2)]]
-var<storage, read> input_2: ArrayVector;
+	[[group(0), binding(2)]]
+	var<storage, read> input_2: ArrayVector;
+
+	[[group(0), binding(3)]]
+	var<storage, write> output_0: Array;
 
-[[group(0), binding(3)]]
-var<storage, write> output_0: Array;
 {%- else -%}
-[[group(0), binding(2)]]
-var<storage, write> output_0: Array;
-{%- endif %}  
+	[[group(0), binding(2)]]
+	var<storage, write> output_0: Array;
+
+{%- endif %}
+
 
-// conv_kernel_1.wgsl
 [[stage(compute), workgroup_size(256, 1, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-        let gidx = global_id.x;
-        if (gidx < {{ o_lens[0]/4 }}u) {
-	let batch = gidx / {{ o_chunks[0][0] / 4 }}u; 
-	let rest = gidx % {{ o_chunks[0][0] / 4 }}u; 
+	let gidx = global_id.x;
+	if (gidx < {{ o_lens[0]/4 }}u) {
+		let batch = gidx / {{ o_chunks[0][0] / 4 }}u; 
+		let rest = gidx % {{ o_chunks[0][0] / 4 }}u; 
 
-        let m = rest / {{ o_chunks[0][1] }}u;
-        let xy = rest % {{ o_chunks[0][1] }}u;
-                
-        var result = vec4<f32>(0.0, 0.0, 0.0, 0.0);
-        
-        let root_index = batch * {{ i_chunks[0][0] }}u + xy;
-        let root_kernel_index = m * {{ channel / 16 * 4 }}u;
+		let m = rest / {{ o_chunks[0][1] }}u;
+		let xy = rest % {{ o_chunks[0][1] }}u;
+			
+		var result = Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0));
+		
+		let root_index = batch * {{ i_chunks[0][0] }}u + xy;
+		let root_kernel_index = m * {{ channel / 16 * 4 }}u;
 
-        for(var c: u32 = 0u; c < {{ channel / 16 }}u; c = c + 1u) {
+		for(var c: u32 = 0u; c < {{ channel / 16 }}u; c = c + 1u) {
+			let base_index = root_index + c * {{ 16 * i_chunks[0][1] }}u;
+			var base_kernel_index = root_kernel_index + c;
+			
+			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];
 
-                let base_index = root_index + c * {{ 16 * i_chunks[0][1] }}u;
-                var base_kernel_index = root_kernel_index + c;
-                
-                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;
 
-                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>(
-                                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>(
-                             matrix_0[index_c_vec],
-                             matrix_1[index_c_vec],
-                             matrix_2[index_c_vec],
-                             matrix_3[index_c_vec],                
-                        ) + result; 
-                }
-	}
-
-        {% if i_lens | length == 3 -%}
-        result = result + input_2.data[m];
-        {%- endif %}
+				let tmp_vec = Vec4(
+					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(
+					matrix_0[index_c_vec],
+					matrix_1[index_c_vec],
+					matrix_2[index_c_vec],
+					matrix_3[index_c_vec],
+				) + result; 
+			}
+		}
 
-{% set activation_input = "result" %}
-{% set activation_output = "result" %}
-{% set activation_type = op_type | replace(from="Conv", to="") %}
-{%- include "snippets/activation_vec.wgsl" %}
+		{% if i_lens | length == 3 -%}
+			result = result + input_2.data[m];
+		{%- endif %}
 
-        let base_index = batch * {{ o_chunks[0][0] }}u + m * {{ o_chunks[0][1] * 4 }}u + xy;
-        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;
+		{% set activation_input = "result" %}
+		{% set activation_output = "result" %}
+		{% set activation_type = op_type | replace(from="Conv", to="") %}
+		{%- include "snippets/activation_vec.wgsl" %}
 
-                output_0.data[index] = result[index_vec];
-        }
-        }
+		let base_index = batch * {{ o_chunks[0][0] }}u + m * {{ o_chunks[0][1] * 4 }}u + xy;
+		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;
+			output_0.data[index] = result[index_vec];
+		}
+	}
 }
diff --git a/wonnx/templates/pool/conv_kernel_3.wgsl b/wonnx/templates/pool/conv_kernel_3.wgsl
index 5f4a365c..d448a721 100644
--- a/wonnx/templates/pool/conv_kernel_3.wgsl
+++ b/wonnx/templates/pool/conv_kernel_3.wgsl
@@ -1,9 +1,5 @@
 {%- include "structs.wgsl" -%}
 
-struct ArrayMatrix3 {
-    data: [[stride(48)]] array<mat3x3<f32>>;
-}; // this is used as both input and output for convenience
-
 [[group(0), binding(0)]]
 var<storage, read> input_0: Array;
 
@@ -11,93 +7,85 @@ var<storage, read> input_0: Array;
 var<storage, read> input_1: ArrayMatrix3;
 
 {%- if i_lens | length == 3 -%} // Bias
-[[group(0), binding(2)]]
-var<storage, read> input_2: ArrayVector;
+	[[group(0), binding(2)]]
+	var<storage, read> input_2: ArrayVector;
 
-[[group(0), binding(3)]]
-var<storage, write> output_0: Array;
+	[[group(0), binding(3)]]
+	var<storage, write> output_0: Array;
 
 {%- else -%}
-[[group(0), binding(2)]]
-var<storage, write> output_0: Array;
-{%- endif %}  
+	[[group(0), binding(2)]]
+	var<storage, write> output_0: Array;
+
+{%- endif %}
 
-// conv_kernel_3.wgsl
 [[stage(compute), workgroup_size(256, 1, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 	let gidx = global_id.x;
-        if (gidx < {{ o_lens[0]/4 }}u) {
-	let batch = gidx / {{ o_chunks[0][0] / 4 }}u; 
-	let rest = gidx % {{ o_chunks[0][0] / 4 }}u; 
-
-        let m = rest / {{ o_chunks[0][1] }}u;
-        let rest = rest % {{ o_chunks[0][1] }}u;
-        
-        let y = rest / {{ o_chunks[0][2] }}u;
-        let x = rest % {{ o_chunks[0][2] }}u;
-        
-        var result = vec4<f32>(0., 0., 0., 0.);
-        
-        let root_index = batch * {{ i_chunks[0][0] }}u;
-        let root_kernel_index = m * {{ channel * 4 }}u;
-
-        for(var c: u32 = 0u; c < {{ channel }}u; c = c + 1u) {
-            
-            let base_index = root_index + c * {{ i_chunks[0][1] }}u;
-            let base_kernel_index = root_kernel_index + c;
-
-            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) {
-
-                var tmp_vec = vec3<f32>(0., 0., 0.);
-
-		let tmp_y = y * {{ stride[0] }}u + i * {{ dilation[0] }}u - {{ pad[0] }}u; 
-
-        	if ((tmp_y < {{ original_height }}u) && (tmp_y >= 0u)) {
-        
-	        for(var j: u32 = 0u; j < {{ kernel_shape[1] }}u; j = j + 1u) { 
-
-                        let tmp_x = x * {{ stride[1] }}u + j * {{ dilation[1] }}u - {{ pad[1] }}u;
-
-                        if ((tmp_x < {{ original_width }}u) && (tmp_x >= 0u)) {
-
-                                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] = input_0.data[tmp_index];
-
-                        }
-  	        }
+	if (gidx < {{ o_lens[0]/4 }}u) {
+		let batch = gidx / {{ o_chunks[0][0] / 4 }}u; 
+		let rest = gidx % {{ o_chunks[0][0] / 4 }}u; 
+
+		let m = rest / {{ o_chunks[0][1] }}u;
+		let rest = rest % {{ o_chunks[0][1] }}u;
+		
+		let y = rest / {{ o_chunks[0][2] }}u;
+		let x = rest % {{ o_chunks[0][2] }}u;
+		
+		var result = Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0));
+		
+		let root_index = batch * {{ i_chunks[0][0] }}u;
+		let root_kernel_index = m * {{ channel * 4 }}u;
+
+		for(var c: u32 = 0u; c < {{ channel }}u; c = c + 1u) {
+			let base_index = root_index + c * {{ i_chunks[0][1] }}u;
+			let base_kernel_index = root_kernel_index + c;
+
+			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) {
+				var tmp_vec = Vec3(Scalar(0), Scalar(0), Scalar(0));
+				let tmp_y = y * {{ stride[0] }}u + i * {{ dilation[0] }}u - {{ pad[0] }}u; 
+				
+				if ((tmp_y < {{ original_height }}u) && (tmp_y >= 0u)) {
+					for(var j: u32 = 0u; j < {{ kernel_shape[1] }}u; j = j + 1u) { 
+						let tmp_x = x * {{ stride[1] }}u + j * {{ dilation[1] }}u - {{ pad[1] }}u;
+
+						if ((tmp_x < {{ original_width }}u) && (tmp_x >= 0u)) {
+							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] = input_0.data[tmp_index];
+						}
+					}
+				}
+
+				result = tmp_vec * Mat4x3(
+					kernel_matrix_0[i],
+					kernel_matrix_1[i],
+					kernel_matrix_2[i],
+					kernel_matrix_3[i]
+				) + result;
+			}
 		}
-
-               result = tmp_vec * mat4x3<f32>(
-                       kernel_matrix_0[i],
-                       kernel_matrix_1[i],
-                       kernel_matrix_2[i],
-                       kernel_matrix_3[i]
-               ) + result;
-            }
-
-        }       
-	
-        {% if i_lens | length == 3 -%}
-        result = result + input_2.data[m];
-        {%- endif %}
-
-{% set activation_input = "result" %}
-{% set activation_output = "result" %}
-{% set activation_type = op_type | replace(from="Conv", to="") %}
-{%- include "snippets/activation_vec.wgsl" -%}
-
-        let base_index = batch * {{ o_chunks[0][0] }}u + m * {{ o_chunks[0][1] * 4 }}u + y * {{ width }}u + x;
-        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;
-
-                output_0.data[index] = result[index_vec];
-        }
-        }
+		
+		{% if i_lens | length == 3 -%}
+			result = result + input_2.data[m];
+		{%- endif %}
+
+		{% set activation_input = "result" %}
+		{% set activation_output = "result" %}
+		{% set activation_type = op_type | replace(from="Conv", to="") %}
+		{%- include "snippets/activation_vec.wgsl" -%}
+
+		let base_index = batch * {{ o_chunks[0][0] }}u + m * {{ o_chunks[0][1] * 4 }}u + y * {{ width }}u + x;
+		
+		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;
+			output_0.data[index] = result[index_vec];
+		}
+	}
 }
diff --git a/wonnx/templates/snippets/activation_scalar.wgsl b/wonnx/templates/snippets/activation_scalar.wgsl
index 1616a1b7..2464ae4d 100644
--- a/wonnx/templates/snippets/activation_scalar.wgsl
+++ b/wonnx/templates/snippets/activation_scalar.wgsl
@@ -1,43 +1,47 @@
 {%- if activation_type == "Relu" -%}
-    {{ activation_output }} = max({{ activation_input }},  0.0);
+	{{ activation_output }} = max({{ activation_input }}, Scalar(0));
 
 {%- elif activation_type == "Sigmoid" -%}
-    {{ activation_output }} = 1.0 / (1.0 + exp(-{{ activation_input }}));
+	{{ activation_output }} = Scalar(1) / (Scalar(1) + exp(-{{ activation_input }}));
 
 {%- elif activation_type == "Softsign" -%}
-    let input = {{ activation_input }}; 
-    {{ activation_output }} = input / (1.0 + abs(input));
+	let input = {{ activation_input }}; 
+	{{ activation_output }} = input / (Scalar(1) + abs(input));
 
 {%- elif activation_type == "Softplus" -%}
-    {{ activation_output }} = log(1.0 + exp({{ activation_input }}));
+	{{ activation_output }} = log(Scalar(1) + exp({{ activation_input }}));
 
 {%- elif activation_type == "Clip" -%}
-    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),
-        vec4<f32>(max_clip, max_clip, max_clip, max_clip),
-    );
+	let min_clip = input_1.data[0u];
+	let max_clip = input_2.data[0u];
+
+	{{ activation_output }} = clamp(
+		{{ activation_input }}, 
+		Vec4(min_clip, min_clip, min_clip, min_clip),
+		Vec4(max_clip, max_clip, max_clip, max_clip),
+	);
 
 {%- elif activation_type == "Celu" -%}
-    let input_vec = {{ activation_input }}; 
-    {{ activation_output }} = max(0.0, 
-            input_vec
-        ) + min(
-            0.0, 
-            {{ alpha }} * (exp(input_vec / {{ alpha }}) - 1.0)
-        );
+	let input_vec = {{ activation_input }};
+
+	{{ activation_output }} = max(Scalar(0), 
+			input_vec
+		) + min(
+			Scalar(0), 
+			{{ alpha }} * (exp(input_vec / {{ alpha }}) - Scalar(1))
+		);
 
 {%- elif activation_type == "Elu" -%}
-        let input_vec = {{ activation_input }}; 
-        {{ activation_output }} = max(
-            0.0, 
-            input_vec
-        ) + min(
-            0.0, 
-            {{ alpha }} * (exp(input_vec) - 1.0)
-        );
+		let input_vec = {{ activation_input }}; 
+		{{ activation_output }} = max(
+			Scalar(0), 
+			input_vec
+		) + min(
+			Scalar(0), 
+			{{ alpha }} * (exp(input_vec) - Scalar(1))
+		);
+
 {%- elif activation_output != activation_input -%}
-       {{ activation_output }} = {{ activation_input }};
+	{{ activation_output }} = {{ activation_input }};
+
 {%- endif -%}
\ No newline at end of file
diff --git a/wonnx/templates/snippets/activation_vec.wgsl b/wonnx/templates/snippets/activation_vec.wgsl
index 8df53aac..4a46730a 100644
--- a/wonnx/templates/snippets/activation_vec.wgsl
+++ b/wonnx/templates/snippets/activation_vec.wgsl
@@ -1,52 +1,53 @@
 {%- if activation_type == "Relu"-%}
-    {{ activation_output }} = max({{ activation_input }}, vec4<f32>(0.0, 0.0, 0.0, 0.0));
+	{{ activation_output }} = max({{ activation_input }}, Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)));
 
 {%- elif activation_type == "Sigmoid" -%}
-    {{ activation_output }} = vec4<f32>(1.0, 1.0, 1.0, 1.0) / (vec4<f32>(1.0, 1.0, 1.0, 1.0) + exp(-{{ activation_input }}));
+	{{ activation_output }} = Vec4(Scalar(1), Scalar(1), Scalar(1), Scalar(1)) / (Vec4(Scalar(1), Scalar(1), Scalar(1), Scalar(1)) + exp(-{{ activation_input }}));
 
 {%- elif activation_type == "Softsign" -%}
-    let input = {{ activation_input }}; 
-    {{ activation_output }} = input / (vec4<f32>(1.0, 1.0, 1.0, 1.0) + abs(input));
+	let input = {{ activation_input }}; 
+	{{ activation_output }} = input / (Vec4(Scalar(1), Scalar(1), Scalar(1), Scalar(1)) + abs(input));
 
 {%- elif activation_type == "Softplus" -%}
-    {{ activation_output }} = log(vec4<f32>(1.0, 1.0, 1.0, 1.0) + exp({{ activation_input }}));
+	{{ activation_output }} = log(Vec4(Scalar(1), Scalar(1), Scalar(1), Scalar(1)) + exp({{ activation_input }}));
 
 {%- elif activation_type == "Clip" -%}
-    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),
-        vec4<f32>(max_clip, max_clip, max_clip, max_clip),
-    );
+	let min_clip = input_1.data[0u];
+	let max_clip = input_2.data[0u];
+	{{ activation_output }} = clamp(
+		{{ activation_input }}, 
+		Vec4(min_clip, min_clip, min_clip, min_clip),
+		Vec4(max_clip, max_clip, max_clip, max_clip),
+	);
 
 {%- elif activation_type == "Celu" -%}
-    let input_vec = {{ activation_input }}; 
-    {{ activation_output }} = max(
-            vec4<f32>(0.0, 0.0, 0.0, 0.0), 
-            input_vec
-        ) + min(
-            vec4<f32>(0.0, 0.0, 0.0, 0.0), 
-            {{ alpha }} * (exp(input_vec / {{ alpha }}) - vec4<f32>(1.0, 1.0, 1.0, 1.0))
-        );
+	let input_vec = {{ activation_input }}; 
+	{{ activation_output }} = max(
+			Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)), 
+			input_vec
+		) + min(
+			Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)), 
+			{{ alpha }} * (exp(input_vec / {{ alpha }}) - Vec4(Scalar(1), Scalar(1), Scalar(1), Scalar(1)))
+		);
 
 {%- elif activation_type == "Elu" -%}
-        let input_vec = {{ activation_input }}; 
-        {{ activation_output }} = max(
-            vec4<f32>(0.0, 0.0, 0.0, 0.0), 
-            input_vec
-        ) + min(
-            vec4<f32>(0.0, 0.0, 0.0, 0.0), 
-            {{ alpha }} * (exp(input_vec) - vec4<f32>(1.0, 1.0, 1.0, 1.0))
-        );
+		let input_vec = {{ activation_input }}; 
+		{{ activation_output }} = max(
+			Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)), 
+			input_vec
+		) + min(
+			Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)), 
+			{{ alpha }} * (exp(input_vec) - Vec4(Scalar(1), Scalar(1), Scalar(1), Scalar(1)))
+		);
 
 {%- elif activation_type == "Mish" -%}
-    let input_vec = {{ activation_input }}; 
-    {{ activation_output }} = input_vec * tanh(log(vec4<f32>(1.0, 1.0, 1.0, 1.0) + exp(input_vec)));
+	let input_vec = {{ activation_input }}; 
+	{{ activation_output }} = input_vec * tanh(log(Vec4(Scalar(1), Scalar(1), Scalar(1), Scalar(1)) + exp(input_vec)));
 
 {%- elif activation_type == "LeakyRelu" -%}
-    {{ activation_output }} = max({{ alpha }} * {{ activation_input }}, vec4<f32>(0.0, 0.0, 0.0, 0.0));
+	{{ activation_output }} = max({{ alpha }} * {{ activation_input }}, Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0)));
 
 {%- elif activation_output != activation_input -%}
-       {{ activation_output }} = {{ activation_input }};
+	{{ activation_output }} = {{ activation_input }};
+
 {%- endif -%}
\ No newline at end of file
diff --git a/wonnx/templates/structs.wgsl b/wonnx/templates/structs.wgsl
index fdb61543..6c7db4d3 100644
--- a/wonnx/templates/structs.wgsl
+++ b/wonnx/templates/structs.wgsl
@@ -1,11 +1,27 @@
+{# 
+// Operations usually work with a single scalar data type (typically f32). This data type is set by the compiler as the
+// 'scalar_type' variable. Here we define several other useful data types that shader code can use to make it more portable.
+#}
+
+type Scalar = {{ scalar_type }};
+type Vec3 = vec3<{{ scalar_type }}>;
+type Vec4 = vec4<{{ scalar_type }}>;
+type Mat3x3 = mat3x3<{{ scalar_type }}>;
+type Mat4x4 = mat4x4<{{ scalar_type }}>;
+type Mat4x3 = mat4x3<{{ scalar_type }}>;
+
 struct Array {
-    data: [[stride(4)]] array<f32>;
+	data: [[stride({{ scalar_stride }})]] array<Scalar>;
 }; 
 
 struct ArrayVector {
-    data: [[stride(16)]] array<vec4<f32>>;
+	data: [[stride({{ vec4_stride }})]] array<Vec4>;
 }; 
 
 struct ArrayMatrix {
-    data: [[stride(64)]] array<mat4x4<f32>>;
+	data: [[stride({{ mat4x4_stride }})]] array<Mat4x4>;
 }; 
+
+struct ArrayMatrix3 {
+	data: [[stride({{ mat3x3_stride }})]] array<Mat3x3>;
+};


Test patch:
diff --git a/wonnx/tests/arithmetic.rs b/wonnx/tests/arithmetic.rs
new file mode 100644
index 00000000..d0c666f1
--- /dev/null
+++ b/wonnx/tests/arithmetic.rs
@@ -0,0 +1,82 @@
+use std::collections::HashMap;
+use wonnx::{
+    onnx::TensorProto_DataType,
+    utils::{graph, model, node, tensor, tensor_of_type, InputTensor},
+};
+
+#[test]
+fn test_cos() {
+    let n: usize = 16;
+    let mut input_data = HashMap::new();
+
+    let data = vec![0.0f32; n];
+    let shape = vec![n as i64];
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
+
+    // Model: X -> Cos -> Y
+    let model = model(graph(
+        vec![tensor("X", &shape)],
+        vec![tensor("Y", &shape)],
+        vec![],
+        vec![],
+        vec![node(vec!["X"], vec!["Y"], "cos", "Cos", 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["Y"], [1.0; 16]);
+}
+
+#[test]
+fn test_reciprocal() {
+    let n: usize = 16;
+    let mut input_data = HashMap::new();
+
+    let data: Vec<f32> = (1..=n).map(|x| x as f32).collect();
+    let reciprocal_data: Vec<f32> = (1..=n).map(|x| 1.0 / (x as f32)).collect();
+    let shape = vec![n as i64];
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
+
+    // Model: X -> Reciprocal -> Y
+    let model = model(graph(
+        vec![tensor("X", &shape)],
+        vec![tensor("Y", &shape)],
+        vec![],
+        vec![],
+        vec![node(vec!["X"], vec!["Y"], "rec", "Reciprocal", 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["Y"], reciprocal_data);
+}
+
+#[test]
+fn test_integer() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let n: usize = 16;
+    let mut input_data = HashMap::new();
+
+    let data = vec![21i32; n];
+    let shape = vec![n as i64];
+    input_data.insert("X".to_string(), InputTensor::I32(data.as_slice()));
+
+    // Model: X -> Cos -> Y
+    let model = model(graph(
+        vec![tensor_of_type("X", &shape, TensorProto_DataType::INT32)],
+        vec![tensor_of_type("Y", &shape, TensorProto_DataType::INT32)],
+        vec![],
+        vec![],
+        vec![node(vec!["X", "X"], vec!["Y"], "add_ints", "Add", 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["Y"], vec![42.0; n]);
+}
diff --git a/wonnx/tests/batchnormalization.rs b/wonnx/tests/batchnormalization.rs
index 02d0f7d0..15546b29 100644
--- a/wonnx/tests/batchnormalization.rs
+++ b/wonnx/tests/batchnormalization.rs
@@ -1,6 +1,6 @@
 use approx::assert_ulps_eq;
 use std::collections::HashMap;
-use wonnx::utils::{attribute, graph, initializer, model, node, tensor};
+use wonnx::utils::{attribute, graph, initializer, model, node, tensor, InputTensor};
 
 fn abs_eq_vector(xs: &[f32], ys: &[f32]) {
     assert_eq!(xs.len(), ys.len());
@@ -25,7 +25,7 @@ fn batch_normalization() {
         width_height as i64,
         width_height as i64,
     ];
-    input_data.insert("X".to_string(), data.as_slice());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     let mean: Vec<f32> = vec![100.0, -100.0];
     let var: Vec<f32> = vec![10.0, 10.0];
diff --git a/wonnx/tests/cast.rs b/wonnx/tests/cast.rs
new file mode 100644
index 00000000..546b6f3e
--- /dev/null
+++ b/wonnx/tests/cast.rs
@@ -0,0 +1,42 @@
+use std::collections::HashMap;
+
+use protobuf::ProtobufEnum;
+use wonnx::{
+    onnx::TensorProto_DataType,
+    utils::{attribute, graph, model, node, tensor, tensor_of_type, InputTensor},
+};
+
+#[test]
+fn test_cast() {
+    let n: usize = 16;
+    let mut input_data = HashMap::new();
+
+    // The input will be 0.0, 0.33, 0.66, ... and (by casting) we'll effectively 'floor' these
+    let data: Vec<f32> = (0..n).map(|x| x as f32 / 3.0).collect();
+    let dims = vec![n as i64];
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
+
+    // Model: X -> Identity -> Y; Y==Z
+    let model = model(graph(
+        vec![tensor("X", &dims)],
+        vec![tensor_of_type("Y", &dims, TensorProto_DataType::INT32)],
+        vec![],
+        vec![],
+        vec![node(
+            vec!["X"],
+            vec!["Y"],
+            "a",
+            "Cast",
+            vec![attribute("to", TensorProto_DataType::INT32.value() as i64)],
+        )],
+    ));
+
+    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["Y"],
+        vec![0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 5.0]
+    );
+}
diff --git a/wonnx/tests/conv.rs b/wonnx/tests/conv.rs
index 5826ad75..bee88dd2 100644
--- a/wonnx/tests/conv.rs
+++ b/wonnx/tests/conv.rs
@@ -2,7 +2,7 @@ use std::collections::HashMap;
 // use wasm_bindgen_test::*;
 use wonnx::*;
 // Indicates a f32 overflow in an intermediate Collatz value
-use wonnx::utils::{attribute, graph, initializer, model, node, tensor};
+use wonnx::utils::{attribute, graph, initializer, model, node, tensor, InputTensor};
 
 async fn run() {
     //    let conv_asymetric_stride = conv_kernel_3().await.unwrap();
@@ -26,7 +26,7 @@ fn conv_pad() {
 
     let data: Vec<f32> = (0..50).map(|x| x as f32).collect();
     let shape = vec![2, c as i64, n as i64, n as i64];
-    input_data.insert("X".to_string(), data.as_slice());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // ONNX INPUTS
 
@@ -81,7 +81,7 @@ fn conv_without_pad() {
 
     let data: Vec<f32> = (0..25).map(|x| x as f32).collect();
     let shape = vec![1, c as i64, n as i64, n as i64];
-    input_data.insert("X".to_string(), data.as_slice());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // ONNX INPUTS
 
@@ -122,7 +122,7 @@ fn conv_stride() {
     let mut input_data = HashMap::new();
 
     let data: Vec<f32> = (0..35).map(|x| x as f32).collect();
-    input_data.insert("X".to_string(), data.as_slice());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // ONNX INPUTS
 
@@ -176,7 +176,7 @@ fn conv_asymetric_stride() {
     let mut input_data = HashMap::new();
 
     let data: Vec<f32> = (0..35).map(|x| x as f32).collect();
-    input_data.insert("X".to_string(), data.as_slice());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // ONNX INPUTS
     let kernel_n = 3;
@@ -222,7 +222,7 @@ fn _conv_kernel_3() {
     let mut input_data = HashMap::new();
 
     let data: Vec<f32> = (0..16).map(|x| x as f32).collect();
-    input_data.insert("X".to_string(), data.as_slice());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // ONNX INPUTS
 
diff --git a/wonnx/tests/cos.rs b/wonnx/tests/cos.rs
deleted file mode 100644
index d219243d..00000000
--- a/wonnx/tests/cos.rs
+++ /dev/null
@@ -1,27 +0,0 @@
-use std::collections::HashMap;
-use wonnx::utils::{graph, model, node, tensor};
-
-#[test]
-fn test_cos() {
-    let n: usize = 16;
-    let mut input_data = HashMap::new();
-
-    let data = vec![0.0f32; n];
-    let shape = vec![n as i64];
-    input_data.insert("X".to_string(), data.as_slice());
-
-    // Model: X -> Cos -> Y
-    let model = model(graph(
-        vec![tensor("X", &shape)],
-        vec![tensor("Y", &shape)],
-        vec![],
-        vec![],
-        vec![node(vec!["X"], vec!["Y"], "cos", "Cos", 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["Y"], [1.0; 16]);
-}
diff --git a/wonnx/tests/globalaveragepool.rs b/wonnx/tests/globalaveragepool.rs
index 714fdcec..9db3ebfd 100644
--- a/wonnx/tests/globalaveragepool.rs
+++ b/wonnx/tests/globalaveragepool.rs
@@ -1,6 +1,6 @@
 use approx::assert_ulps_eq;
 use std::collections::HashMap;
-use wonnx::utils::{graph, model, node, tensor};
+use wonnx::utils::{graph, model, node, tensor, InputTensor};
 
 fn abs_eq_vector(xs: &[f32], ys: &[f32]) {
     assert_eq!(xs.len(), ys.len());
@@ -27,7 +27,7 @@ fn global_average_pool() {
         width_height as i64,
     ];
     let output_shape = vec![batches as i64, channels as i64, 1, 1];
-    input_data.insert("X".to_string(), data.as_slice());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // Model: X -> GlobalAveragePool -> Y
     let bn_model = model(graph(
diff --git a/wonnx/tests/identity.rs b/wonnx/tests/identity.rs
index ba15460b..0982f93b 100644
--- a/wonnx/tests/identity.rs
+++ b/wonnx/tests/identity.rs
@@ -2,7 +2,7 @@ use std::collections::HashMap;
 // use wasm_bindgen_test::*;
 // Indicates a f32 overflow in an intermediate Collatz value
 
-use wonnx::utils::{graph, model, node, tensor};
+use wonnx::utils::{graph, model, node, tensor, InputTensor};
 
 #[test]
 fn test_identity() {
@@ -11,7 +11,7 @@ fn test_identity() {
 
     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());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // Model: X -> Identity -> Y; Y==Z
     let model = model(graph(
@@ -36,7 +36,7 @@ fn test_double_identity() {
 
     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());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // Model: X -> Identity -> Y -> Identity -> Z. X==Z
     let model = model(graph(
diff --git a/wonnx/tests/matrix.rs b/wonnx/tests/matrix.rs
index fa6e82f7..d63bfebd 100644
--- a/wonnx/tests/matrix.rs
+++ b/wonnx/tests/matrix.rs
@@ -1,5 +1,5 @@
 use std::collections::HashMap;
-use wonnx::utils::{attribute, graph, initializer, model, node, tensor};
+use wonnx::utils::{attribute, graph, initializer, model, node, tensor, InputTensor};
 
 #[test]
 fn test_matmul_square_matrix() {
@@ -20,8 +20,14 @@ fn test_matmul_square_matrix() {
 
     let sum = data_a.dot(&data_b);
 
-    input_data.insert("A".to_string(), data_a.as_slice().unwrap());
-    input_data.insert("B".to_string(), data_b.as_slice().unwrap());
+    input_data.insert(
+        "A".to_string(),
+        InputTensor::F32(data_a.as_slice().unwrap()),
+    );
+    input_data.insert(
+        "B".to_string(),
+        InputTensor::F32(data_b.as_slice().unwrap()),
+    );
 
     let n = n as i64;
     let model = model(graph(
@@ -44,7 +50,7 @@ fn test_matmul_square_matrix() {
 fn test_two_transposes() {
     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());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     // Model: X -> Transpose -> Y -> Transpose -> Z; X==Z
     let model = model(graph(
@@ -81,7 +87,7 @@ fn test_two_transposes() {
 fn test_split() {
     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());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     let model = model(graph(
         vec![tensor("X", &[2, 6])],
@@ -112,7 +118,7 @@ fn test_resize() {
     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());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     let downsampling_model = model(graph(
         vec![tensor("X", &[1, 1, 2, 4])],
@@ -137,7 +143,7 @@ fn test_resize() {
 
     let mut input_data = HashMap::new();
     let data = (1..=4).map(|x| x as f32).collect::<Vec<f32>>();
-    input_data.insert("X".to_string(), data.as_slice());
+    input_data.insert("X".to_string(), InputTensor::F32(data.as_slice()));
 
     let upsampling_model = model(graph(
         vec![tensor("X", &[1, 1, 2, 2])],
@@ -145,7 +151,7 @@ fn test_resize() {
         vec![],
         vec![initializer("scales", vec![1., 1., 2., 3.])],
         vec![node(
-            vec!["X", "scales"],
+            vec!["X", "" /* roi */, "scales"],
             vec!["Y"],
             "Resize",
             "Resize",
diff --git a/wonnx/tests/pretrained_models.rs b/wonnx/tests/pretrained_models.rs
index 24af1f03..187a8f9b 100644
--- a/wonnx/tests/pretrained_models.rs
+++ b/wonnx/tests/pretrained_models.rs
@@ -4,12 +4,13 @@ use std::collections::HashMap;
 use std::path::Path;
 #[cfg(target_arch = "wasm32")]
 use wasm_bindgen_test::*;
+use wonnx::utils::InputTensor;
 
 #[test]
 fn test_relu() {
     let mut input_data = HashMap::new();
     let data = vec![-1.0f32, 1.0];
-    input_data.insert("x".to_string(), data.as_slice());
+    input_data.insert("x".to_string(), InputTensor::F32(data.as_slice()));
 
     let session = pollster::block_on(wonnx::Session::from_path(
         "examples/data/models/single_relu.onnx",
@@ -25,7 +26,10 @@ 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());
+    input_data.insert(
+        "Input3".to_string(),
+        InputTensor::F32(image.as_slice().unwrap()),
+    );
     let session = pollster::block_on(wonnx::Session::from_path(
         "examples/data/models/opt-mnist.onnx",
     ))
@@ -46,7 +50,10 @@ 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());
+    input_data.insert(
+        "Input3".to_string(),
+        InputTensor::F32(image.as_slice().unwrap()),
+    );
     let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
         .iter()
         .enumerate()
@@ -62,7 +69,10 @@ 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());
+    input_data.insert(
+        "Input3".to_string(),
+        InputTensor::F32(image.as_slice().unwrap()),
+    );
     let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
         .iter()
         .enumerate()
@@ -78,7 +88,10 @@ 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());
+    input_data.insert(
+        "Input3".to_string(),
+        InputTensor::F32(image.as_slice().unwrap()),
+    );
     let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
         .iter()
         .enumerate()
@@ -100,7 +113,10 @@ fn test_mnist() {
 fn test_squeeze() {
     let mut input_data = HashMap::new();
     let image = load_squeezenet_image();
-    input_data.insert("data".to_string(), image.as_slice().unwrap());
+    input_data.insert(
+        "data".to_string(),
+        InputTensor::F32(image.as_slice().unwrap()),
+    );
 
     let session = pollster::block_on(wonnx::Session::from_path(
         "examples/data/models/opt-squeeze.onnx",