Repository: webonnx/wonnx (language: python)

Issue:
Process i64 data as i32 data
WGSL does not appear to support 64-bit integers, so the best we can do (apart from building all sorts of complicated 64-bit logic based on 32 bit ints) is to use i32 instead (this also means: load i64 tensors as i32) and hope for the best (and wait for WGSL to start supporting 64 bit ints).

In practice it appears models use i64 but do not need the full range (i.e. in one of the ONNX files for BERTSQuAD, the input tokens are i64 tensors, but these values would also fit in i32).

Still need to implement the 64-to-32 conversion when loading tensors.


Discussion / hints:
Current version supports reading int64-values inputs and initializers as int32 values (which is transparent as long as all the values fit into an int32...). See `test_int64_initializers` in `wonnx/tests/arithmetic.rs` for an example. Wonnx will log the appropriate warnings. This is supposed to be a temporary hack until WGSL supports int64 as data type (at that point we can implement support without changing the user-facing API). 
After this is merged I should do some more work to also support non-f64 outputs (including i64 outputs which are i32 outputs internally)
@haixuanTao OK so I have now made the conversions fallible (note that `From<i64>` is not implemented for `f32` so `i64.try_into()` will not work, instead I use `f32::from_i64` from `num`.

Fix patch:
diff --git a/README.md b/README.md
index 9fb0430c..f9cce27c 100644
--- a/README.md
+++ b/README.md
@@ -195,7 +195,9 @@ fn test_matmul_square_matrix() {
 
     let result = pollster::block_on(session.run(input_data)).unwrap();
 
-    assert_eq!(result["C"].as_slice(), sum.as_slice().unwrap());
+    // Note: it is better to use a method that compares floats with a tolerance to account for differences
+    // between implementations; see `wonnx/tests/common/mod.rs` for an example.
+    assert_eq!((&result["C"]).try_into().unwrap(),sum.as_slice().unwrap());
 }
 ```
 > Check out tera documentation for other templating operation: https://tera.netlify.app/docs/
diff --git a/wonnx-cli/src/cpu.rs b/wonnx-cli/src/cpu.rs
index 6818aefa..0a936e61 100644
--- a/wonnx-cli/src/cpu.rs
+++ b/wonnx-cli/src/cpu.rs
@@ -3,7 +3,10 @@ use std::collections::HashMap;
 use crate::{InferOptions, Inferer, NNXError};
 use async_trait::async_trait;
 use tract_onnx::prelude::*;
-use wonnx::{onnx::ModelProto, utils::Shape};
+use wonnx::{
+    onnx::ModelProto,
+    utils::{OutputTensor, Shape},
+};
 
 type RunnableOnnxModel =
     SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>;
@@ -62,7 +65,7 @@ impl Inferer for CPUInferer {
         infer_opt: &InferOptions,
         inputs: &HashMap<String, crate::Tensor>,
         model: &ModelProto,
-    ) -> Result<Vec<f32>, NNXError> {
+    ) -> Result<OutputTensor, NNXError> {
         let mut cpu_inputs: HashMap<usize, tract_onnx::prelude::Tensor> = HashMap::new();
 
         for (input_name, input_tensor) in inputs {
@@ -119,6 +122,6 @@ impl Inferer for CPUInferer {
         };
 
         let av = result_vector.to_array_view()?;
-        Ok(av.as_slice().unwrap().to_vec())
+        Ok(OutputTensor::F32(av.as_slice().unwrap().to_vec()))
     }
 }
diff --git a/wonnx-cli/src/gpu.rs b/wonnx-cli/src/gpu.rs
index 5922f67d..d1bf3928 100644
--- a/wonnx-cli/src/gpu.rs
+++ b/wonnx-cli/src/gpu.rs
@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use wonnx::onnx::ModelProto;
 
 use async_trait::async_trait;
+use wonnx::utils::OutputTensor;
 
 use crate::types::InferOptions;
 use crate::types::Inferer;
@@ -27,7 +28,7 @@ impl Inferer for GPUInferer {
         infer_opt: &InferOptions,
         inputs: &HashMap<String, crate::Tensor>,
         _model: &ModelProto,
-    ) -> Result<Vec<f32>, NNXError> {
+    ) -> Result<OutputTensor, NNXError> {
         let input_refs = inputs
             .iter()
             .map(|(k, v)| (k.clone(), v.input_tensor()))
diff --git a/wonnx-cli/src/main.rs b/wonnx-cli/src/main.rs
index 496da083..37845aab 100644
--- a/wonnx-cli/src/main.rs
+++ b/wonnx-cli/src/main.rs
@@ -159,7 +159,10 @@ async fn run() -> Result<(), NNXError> {
                         let _ = gpu_backend.infer(&infer_opt, &inputs, &model).await?;
                     }
                 }
-                let gpu_output = gpu_backend.infer(&infer_opt, &inputs, &model).await?;
+                let gpu_output: Vec<f32> = gpu_backend
+                    .infer(&infer_opt, &inputs, &model)
+                    .await?
+                    .try_into()?;
                 let gpu_time = gpu_start.elapsed();
                 log::info!("gpu time: {}ms", gpu_time.as_millis());
                 drop(gpu_backend);
@@ -173,7 +176,10 @@ async fn run() -> Result<(), NNXError> {
                         let _ = cpu_backend.infer(&infer_opt, &inputs, &model).await?;
                     }
                 }
-                let cpu_output = cpu_backend.infer(&infer_opt, &inputs, &model).await?;
+                let cpu_output: Vec<f32> = cpu_backend
+                    .infer(&infer_opt, &inputs, &model)
+                    .await?
+                    .try_into()?;
                 let cpu_time = cpu_start.elapsed();
                 log::info!(
                     "cpu time: {}ms ({:.2}x gpu time)",
@@ -265,7 +271,8 @@ async fn run() -> Result<(), NNXError> {
                 Some(labels_path) => {
                     let labels = get_lines(&labels_path);
 
-                    let mut probabilities = output.iter().enumerate().collect::<Vec<_>>();
+                    let output_slice: Vec<f32> = output.try_into().unwrap();
+                    let mut probabilities = output_slice.iter().enumerate().collect::<Vec<_>>();
                     probabilities.sort_unstable_by(|a, b| b.1.partial_cmp(a.1).unwrap());
 
                     let top = infer_opt.top.unwrap_or(10);
diff --git a/wonnx-cli/src/types.rs b/wonnx-cli/src/types.rs
index a30c5bb5..cb70f657 100644
--- a/wonnx-cli/src/types.rs
+++ b/wonnx-cli/src/types.rs
@@ -1,7 +1,11 @@
 use std::{collections::HashMap, num::ParseFloatError, path::PathBuf, str::FromStr};
 use structopt::StructOpt;
 use thiserror::Error;
-use wonnx::{onnx::ModelProto, SessionError, WonnxError};
+use wonnx::{
+    onnx::ModelProto,
+    utils::{OutputTensor, TensorConversionError},
+    SessionError, WonnxError,
+};
 use wonnx_preprocessing::{text::PreprocessingError, Tensor};
 
 #[cfg(feature = "cpu")]
@@ -54,6 +58,9 @@ pub enum NNXError {
 
     #[error("invalid number: {0}")]
     InvalidNumber(ParseFloatError),
+
+    #[error("tensor error: {0}")]
+    TensorConversionError(#[from] TensorConversionError),
 }
 
 impl FromStr for Backend {
@@ -181,5 +188,5 @@ pub trait Inferer {
         infer_opt: &InferOptions,
         inputs: &HashMap<String, Tensor>,
         model: &ModelProto,
-    ) -> Result<Vec<f32>, NNXError>;
+    ) -> Result<OutputTensor, NNXError>;
 }
diff --git a/wonnx-py/src/lib.rs b/wonnx-py/src/lib.rs
index c807b082..6e2d9ead 100644
--- a/wonnx-py/src/lib.rs
+++ b/wonnx-py/src/lib.rs
@@ -1,6 +1,7 @@
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
 use std::collections::HashMap;
+use wonnx::utils::OutputTensor;
 
 use wonnx::Session;
 #[pyclass]
@@ -9,6 +10,18 @@ pub struct PySession {
     pub session: Session,
 }
 
+pub struct PyOutputTensor(OutputTensor);
+
+impl IntoPy<PyObject> for PyOutputTensor {
+    fn into_py(self, py: Python) -> PyObject {
+        match self.0 {
+            OutputTensor::F32(fs) => fs.into_py(py),
+            OutputTensor::I32(fs) => fs.into_py(py),
+            OutputTensor::I64(fs) => fs.into_py(py),
+        }
+    }
+}
+
 #[pymethods]
 impl PySession {
     #[staticmethod]
@@ -23,15 +36,17 @@ impl PySession {
         Ok(PySession { session })
     }
 
-    pub fn run(&self, dict: &PyDict) -> PyResult<HashMap<String, Vec<f32>>> {
+    pub fn run(&self, dict: &PyDict) -> PyResult<HashMap<String, PyOutputTensor>> {
         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().into());
         }
         let result = pollster::block_on(self.session.run(&inputs)).unwrap();
-
-        Ok(result)
+        Ok(result
+            .into_iter()
+            .map(|(k, v)| (k, PyOutputTensor(v)))
+            .collect())
     }
 }
 
diff --git a/wonnx-wasm/src/lib.rs b/wonnx-wasm/src/lib.rs
index e3d4af21..dc41d66d 100644
--- a/wonnx-wasm/src/lib.rs
+++ b/wonnx-wasm/src/lib.rs
@@ -8,7 +8,7 @@ use std::sync::Arc;
 use wasm_bindgen::prelude::*;
 use wasm_bindgen_console_logger::DEFAULT_LOGGER;
 use wasm_bindgen_futures::future_to_promise;
-use wonnx::utils::InputTensor;
+use wonnx::utils::{InputTensor, OutputTensor};
 
 #[wasm_bindgen(start)]
 pub fn main() {
@@ -90,3 +90,12 @@ impl Session {
         })
     }
 }
+
+/// Convert an OutputTensor to a JsValue (we cannot implement Into<JsValue> for OutputTensor here)
+pub fn tensor_to_js_value(tensor: OutputTensor) -> JsValue {
+    match tensor {
+        OutputTensor::F32(fs) => JsValue::from_serde(&fs).unwrap(),
+        OutputTensor::I32(ints) => JsValue::from_serde(&ints).unwrap(),
+        OutputTensor::I64(ints) => JsValue::from_serde(&ints).unwrap(),
+    }
+}
diff --git a/wonnx/Cargo.toml b/wonnx/Cargo.toml
index 3d3b0ab9..c3b0230a 100644
--- a/wonnx/Cargo.toml
+++ b/wonnx/Cargo.toml
@@ -27,7 +27,7 @@ tera = { version = "1.15.0", default-features = false }
 lazy_static = "1.4.0"
 thiserror = "^1.0.0"
 serde_derive = "1.0.133"
-serde = "1.0.133"
+serde = {version = "1.0.133", features = ["derive"] }
 num = "0.4.0"
 
 [dev-dependencies]
diff --git a/wonnx/examples/mnist.rs b/wonnx/examples/mnist.rs
index 65a93f24..269a1ad5 100644
--- a/wonnx/examples/mnist.rs
+++ b/wonnx/examples/mnist.rs
@@ -1,13 +1,16 @@
 use std::collections::HashMap;
+use std::convert::TryInto;
 
 use image::{imageops::FilterType, ImageBuffer, Pixel, Rgb};
 use std::path::Path;
 use std::time::Instant;
+use wonnx::utils::OutputTensor;
 
 // Args Management
 async fn run() {
     let probabilities = execute_gpu().await.unwrap();
-    let (_, probabilities) = probabilities.iter().next().unwrap();
+    let (_, probabilities) = probabilities.into_iter().next().unwrap();
+    let probabilities: Vec<f32> = probabilities.try_into().unwrap();
     println!("steps: {:#?}", probabilities);
     println!("steps: {:#?}", probabilities.len());
 
@@ -19,7 +22,7 @@ async fn run() {
 }
 
 // Hardware management
-async fn execute_gpu() -> Option<HashMap<String, Vec<f32>>> {
+async fn execute_gpu() -> Option<HashMap<String, OutputTensor>> {
     let mut input_data = HashMap::new();
 
     let image = load_image();
diff --git a/wonnx/examples/simple_graph.rs b/wonnx/examples/simple_graph.rs
index e42d2207..12d55b9f 100644
--- a/wonnx/examples/simple_graph.rs
+++ b/wonnx/examples/simple_graph.rs
@@ -1,20 +1,23 @@
 use std::collections::HashMap;
 
 use wonnx::{
-    utils::{attribute, graph, initializer, model, node, tensor},
+    utils::{attribute, graph, initializer, model, node, tensor, OutputTensor},
     SessionError, WonnxError,
 };
 
 async fn run() -> Result<(), WonnxError> {
-    let result = &execute_gpu().await?;
-    let (_, result) = result.iter().next().unwrap();
+    let result = execute_gpu().await?;
+    let result = result.into_iter().next().unwrap().1;
 
-    assert_eq!(result, &[54., 63., 72., 99., 108., 117., 144., 153., 162.]);
+    assert_eq!(
+        result,
+        OutputTensor::F32(vec![54., 63., 72., 99., 108., 117., 144., 153., 162.])
+    );
     Ok(())
 }
 
 // Hardware management
-async fn execute_gpu() -> Result<HashMap<String, Vec<f32>>, SessionError> {
+async fn execute_gpu() -> Result<HashMap<String, OutputTensor>, SessionError> {
     // USER INPUT
     let n = 5;
     let c = 1;
diff --git a/wonnx/examples/squeeze.rs b/wonnx/examples/squeeze.rs
index 37b181a3..159a184c 100644
--- a/wonnx/examples/squeeze.rs
+++ b/wonnx/examples/squeeze.rs
@@ -2,18 +2,21 @@ use image::{imageops::FilterType, ImageBuffer, Pixel, Rgb};
 use log::info;
 use ndarray::s;
 use std::collections::HashMap;
+use std::convert::TryInto;
 use std::time::Instant;
 use std::{
     fs,
     io::{BufRead, BufReader},
     path::Path,
 };
+use wonnx::utils::OutputTensor;
 use wonnx::WonnxError;
 
 // Args Management
 async fn run() {
-    let probabilities = &execute_gpu().await.unwrap();
-    let (_, probabilities) = probabilities.iter().next().unwrap();
+    let probabilities = execute_gpu().await.unwrap();
+    let probabilities = probabilities.into_iter().next().unwrap().1;
+    let probabilities: Vec<f32> = probabilities.try_into().unwrap();
     let mut probabilities = probabilities.iter().enumerate().collect::<Vec<_>>();
     probabilities.sort_unstable_by(|a, b| b.1.partial_cmp(a.1).unwrap());
 
@@ -28,7 +31,7 @@ async fn run() {
 }
 
 // Hardware management
-async fn execute_gpu() -> Result<HashMap<String, Vec<f32>>, WonnxError> {
+async fn execute_gpu() -> Result<HashMap<String, OutputTensor>, WonnxError> {
     let mut input_data = HashMap::new();
     let image = load_image();
     input_data.insert("data".to_string(), image.as_slice().unwrap().into());
diff --git a/wonnx/src/gpu.rs b/wonnx/src/gpu.rs
index 0dbbd64b..4f5e4fc2 100644
--- a/wonnx/src/gpu.rs
+++ b/wonnx/src/gpu.rs
@@ -14,7 +14,10 @@ use crate::{
     ir::{Node, NodeDefinition, NodeIdentifier, OperatorDefinition},
     onnx::TensorProto,
     resource::{self, resize},
-    utils::{ceil, DataTypeError, InputTensor, ScalarType, Shape, MINIMUM_BUFFER_SIZE_BYTES},
+    utils::{
+        ceil, DataTypeError, InputTensor, OutputTensor, ScalarType, Shape,
+        MINIMUM_BUFFER_SIZE_BYTES,
+    },
 };
 
 /// The maximum number of bindings in a binding group (defined by wgpu)
@@ -292,7 +295,7 @@ impl GpuModel {
     pub async fn infer<'a>(
         &self,
         inference_inputs: &HashMap<String, InputTensor<'a>>,
-    ) -> Result<HashMap<String, Vec<f32>>, GpuError> {
+    ) -> Result<HashMap<String, OutputTensor>, GpuError> {
         log::info!("encode inference steps");
         let mut encoder = self
             .device
@@ -310,25 +313,15 @@ impl GpuModel {
     async fn read_outputs<'a>(
         &self,
         inference_inputs: &HashMap<String, InputTensor<'a>>,
-    ) -> Result<HashMap<String, Vec<f32>>, GpuError> {
-        let mut output_data = HashMap::new();
+    ) -> Result<HashMap<String, OutputTensor>, GpuError> {
+        let mut output_data: HashMap<String, OutputTensor> = HashMap::new();
 
         for (output_name, output_source) in &self.inference_outputs {
             output_data.insert(
                 output_name.to_string(),
                 match output_source {
                     InferenceOutput::InferenceInput(input_name) => {
-                        match &inference_inputs[input_name] {
-                            InputTensor::F32(v) => v.to_vec(),
-                            InputTensor::I32(v) => v
-                                .iter()
-                                .map(|i| f32::from_i32(*i).ok_or(GpuError::OutOfBoundsError))
-                                .collect::<Result<Vec<f32>, _>>()?,
-                            InputTensor::I64(v) => v
-                                .iter()
-                                .map(|i| f32::from_i64(*i).ok_or(GpuError::OutOfBoundsError))
-                                .collect::<Result<Vec<f32>, _>>()?,
-                        }
+                        (&inference_inputs[input_name]).into()
                     }
                     InferenceOutput::Tensor(tensor) => {
                         tensor.read_to_vec(&self.device, &self.queue).await?
@@ -648,7 +641,7 @@ impl GpuTensor {
         &self,
         device: &wgpu::Device,
         queue: &wgpu::Queue,
-    ) -> Result<Vec<f32>, GpuError> {
+    ) -> Result<OutputTensor, GpuError> {
         let buffer_slice = self.buffer.slice(..);
 
         // On wgpu we can MAP_READ a buffer that is also used as STORAGE, but WebGPU (on at least Chrome)
@@ -672,17 +665,17 @@ impl GpuTensor {
         // 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 = match self.shape.data_type {
-            ScalarType::F32 => bytemuck::cast_slice(&output_data)[..output_buffer_size].to_vec(),
+            ScalarType::F32 => {
+                OutputTensor::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()
+                OutputTensor::I32(bytemuck::cast_slice(&output_data)[..output_buffer_size].to_vec())
             }
             ScalarType::I64 => {
                 log::warn!("reading int64 output as int32 because internally int64 scalars are not supported");
                 let result_ints: Vec<i32> =
                     bytemuck::cast_slice(&output_data)[..output_buffer_size].to_vec();
-                result_ints.iter().map(|i| *i as f32).collect()
+                OutputTensor::I64(result_ints.iter().map(|i| *i as i64).collect())
             }
         };
         drop(output_data);
diff --git a/wonnx/src/lib.rs b/wonnx/src/lib.rs
index af84215c..dba051d8 100644
--- a/wonnx/src/lib.rs
+++ b/wonnx/src/lib.rs
@@ -17,7 +17,7 @@ use protobuf::{self, Message, ProtobufError};
 use std::collections::HashMap;
 use std::path::Path;
 use std::result::Result;
-use utils::{DataTypeError, InputTensor};
+use utils::{DataTypeError, InputTensor, OutputTensor};
 
 use crate::gpu::GpuModel;
 use thiserror::Error;
@@ -144,7 +144,7 @@ impl Session {
     pub async fn run<'a>(
         &self,
         inputs: &HashMap<String, InputTensor<'a>>,
-    ) -> Result<HashMap<String, Vec<f32>>, SessionError> {
+    ) -> Result<HashMap<String, OutputTensor>, SessionError> {
         Ok(self.gpu_model.infer(inputs).await?)
     }
 }
diff --git a/wonnx/src/utils.rs b/wonnx/src/utils.rs
index a32e0b98..c94ee22f 100644
--- a/wonnx/src/utils.rs
+++ b/wonnx/src/utils.rs
@@ -1,13 +1,16 @@
 use protobuf::ProtobufEnum;
 use protobuf::RepeatedField;
+use serde::Serialize;
 
 use crate::onnx;
 use crate::onnx::OperatorSetIdProto;
 use crate::onnx::TensorProto_DataType;
 use crate::onnx::ValueInfoProto;
+use num::FromPrimitive;
 use std::borrow::Cow;
 use std::convert::From;
 use std::convert::Into;
+use std::convert::TryFrom;
 use std::fmt::Display;
 use std::str::from_utf8;
 use thiserror::Error;
@@ -86,6 +89,67 @@ impl<'a> From<&'a [i64]> for InputTensor<'a> {
     }
 }
 
+#[derive(Error, Debug)]
+pub enum TensorConversionError {
+    #[error("could not convert to the requested type becaue a value could not be represented in the target type")]
+    OutOfBoundsError,
+
+    #[error("cold not return the requested type; conversions cannot be done for slices")]
+    DataTypeError,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize)]
+#[serde(untagged)]
+pub enum OutputTensor {
+    F32(Vec<f32>),
+    I32(Vec<i32>),
+    I64(Vec<i64>),
+}
+
+impl TryFrom<OutputTensor> for Vec<f32> {
+    type Error = TensorConversionError;
+
+    /// Convert OutputTensor into a Vec<f32>, possibly converting integer tensors if the values fit
+    fn try_from(value: OutputTensor) -> Result<Self, Self::Error> {
+        match value {
+            OutputTensor::F32(floats) => Ok(floats),
+            OutputTensor::I32(ints) => ints
+                .into_iter()
+                .map(|i| f32::from_i32(i).ok_or(TensorConversionError::OutOfBoundsError))
+                .collect::<Result<_, _>>(),
+            OutputTensor::I64(ints) => ints
+                .into_iter()
+                .map(|i| f32::from_i64(i).ok_or(TensorConversionError::OutOfBoundsError))
+                .collect::<Result<_, _>>(),
+        }
+    }
+}
+
+/// Convert &OutputTensor into an &[f32]. Because we cannot store converted results, this operation does not attempt
+/// to convert the tensor if the values are of a different type
+impl<'a> TryFrom<&'a OutputTensor> for &'a [f32] {
+    type Error = TensorConversionError;
+
+    fn try_from(value: &'a OutputTensor) -> Result<Self, Self::Error> {
+        match value {
+            OutputTensor::F32(floats) => Ok(floats.as_slice()),
+            OutputTensor::I32(_) | OutputTensor::I64(_) => {
+                Err(TensorConversionError::DataTypeError)
+            }
+        }
+    }
+}
+
+impl<'a> From<&InputTensor<'a>> for OutputTensor {
+    fn from(input: &InputTensor<'a>) -> Self {
+        match input {
+            InputTensor::F32(fs) => OutputTensor::F32(fs.to_vec()),
+            InputTensor::I32(fs) => OutputTensor::I32(fs.to_vec()),
+            InputTensor::I64(fs) => OutputTensor::I64(fs.to_vec()),
+        }
+    }
+}
+
 #[derive(Error, Debug)]
 pub enum DataTypeError {
     #[error("the ONNX scalar data type '{0:?}' is not supported")]
@@ -466,7 +530,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, OutputTensor};
 
     #[test]
     fn test_model() {
@@ -507,7 +571,7 @@ mod tests {
 
         assert_eq!(
             result["Y"],
-            [54., 63., 72., 99., 108., 117., 144., 153., 162.]
+            OutputTensor::F32(vec![54., 63., 72., 99., 108., 117., 144., 153., 162.])
         );
     }
 }


Test patch:
diff --git a/wonnx/tests/arithmetic.rs b/wonnx/tests/arithmetic.rs
index ee7ffd7d..bef525ba 100644
--- a/wonnx/tests/arithmetic.rs
+++ b/wonnx/tests/arithmetic.rs
@@ -1,7 +1,9 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, convert::TryInto};
 use wonnx::{
     onnx::TensorProto_DataType,
-    utils::{graph, initializer_int64, model, node, tensor, tensor_of_type, InputTensor},
+    utils::{
+        graph, initializer_int64, model, node, tensor, tensor_of_type, InputTensor, OutputTensor,
+    },
 };
 
 mod common;
@@ -28,7 +30,7 @@ fn test_cos() {
         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]);
+    assert_eq!(result["Y"], OutputTensor::F32(vec![1.0; 16]));
 }
 
 #[test]
@@ -54,7 +56,10 @@ fn test_reciprocal() {
         pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
 
     let result = pollster::block_on(session.run(&input_data)).unwrap();
-    common::assert_eq_vector(result["Y"].as_slice(), &reciprocal_data);
+    common::assert_eq_vector(
+        (&result["Y"]).try_into().unwrap(),
+        reciprocal_data.as_slice(),
+    );
 }
 
 #[test]
@@ -80,7 +85,7 @@ fn test_integer() {
         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]);
+    assert_eq!(result["Y"], OutputTensor::I32(vec![42; n]));
 }
 
 #[test]
@@ -89,7 +94,7 @@ fn test_int64_initializers() {
     let n: usize = 16;
     let left: Vec<i64> = (0..n).map(|x| x as i64).collect();
     let right: Vec<i64> = (0..n).map(|x| (x * 2) as i64).collect();
-    let sum: Vec<f32> = (0..n).map(|x| (x * 3) as f32).collect();
+    let sum: Vec<i64> = (0..n).map(|x| (x * 3) as i64).collect();
     let dims = vec![n as i64];
 
     let model = model(graph(
@@ -107,5 +112,5 @@ fn test_int64_initializers() {
     input_data.insert("X".to_string(), left.as_slice().into());
     let result = pollster::block_on(session.run(&input_data)).unwrap();
 
-    common::assert_eq_vector(result["Z"].as_slice(), &sum);
+    assert_eq!(result["Z"], OutputTensor::I64(sum))
 }
diff --git a/wonnx/tests/batchnormalization.rs b/wonnx/tests/batchnormalization.rs
index 45d6fd2e..35180150 100644
--- a/wonnx/tests/batchnormalization.rs
+++ b/wonnx/tests/batchnormalization.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, convert::TryInto};
 use wonnx::utils::{attribute, graph, initializer, model, node, tensor};
 mod common;
 
@@ -62,7 +62,7 @@ fn batch_normalization() {
     let out_y = &result["Y"];
 
     common::assert_eq_vector(
-        out_y.as_slice(),
+        out_y.try_into().unwrap(),
         &[
             // Y = (X - input_mean) / sqrt(input_var + epsilon) * scale + B
             // For X=0, Y = (0 - 100) / sqrt(10 + 0.1) * 1 + 1000 = 968,53
diff --git a/wonnx/tests/cast.rs b/wonnx/tests/cast.rs
index 9047ad65..d7948dc3 100644
--- a/wonnx/tests/cast.rs
+++ b/wonnx/tests/cast.rs
@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use protobuf::ProtobufEnum;
 use wonnx::{
     onnx::TensorProto_DataType,
-    utils::{attribute, graph, model, node, tensor, tensor_of_type},
+    utils::{attribute, graph, model, node, tensor, tensor_of_type, OutputTensor},
 };
 
 #[test]
@@ -37,6 +37,6 @@ fn test_cast() {
     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]
+        OutputTensor::I32(vec![0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5])
     );
 }
diff --git a/wonnx/tests/conv.rs b/wonnx/tests/conv.rs
index c8459f10..7c7cf8f6 100644
--- a/wonnx/tests/conv.rs
+++ b/wonnx/tests/conv.rs
@@ -1,6 +1,8 @@
 use std::collections::HashMap;
-use wonnx::utils::{attribute, graph, initializer, model, node, tensor};
+use std::convert::TryInto;
+use wonnx::utils::{attribute, graph, initializer, model, node, tensor, OutputTensor};
 use wonnx::*;
+mod common;
 
 #[test]
 fn conv_pad() {
@@ -36,7 +38,7 @@ fn conv_pad() {
     let result = pollster::block_on(session.run(&input_data)).unwrap();
     assert_eq!(
         result["Y"],
-        [
+        OutputTensor::F32(vec![
             12.0, 21.0, 27.0, 33.0, 24.0, 33.0, 54.0, 63.0, 72.0, 51.0, 63.0, 99.0, 108.0, 117.0,
             81.0, 93.0, 144.0, 153.0, 162.0, 111.0, 72.0, 111.0, 117.0, 123.0, 84.0, 12.0, 21.0,
             27.0, 33.0, 24.0, 33.0, 54.0, 63.0, 72.0, 51.0, 63.0, 99.0, 108.0, 117.0, 81.0, 93.0,
@@ -45,7 +47,7 @@ fn conv_pad() {
             243.0, 369.0, 378.0, 387.0, 261.0, 172.0, 261.0, 267.0, 273.0, 184.0, 112.0, 171.0,
             177.0, 183.0, 124.0, 183.0, 279.0, 288.0, 297.0, 201.0, 213.0, 324.0, 333.0, 342.0,
             231.0, 243.0, 369.0, 378.0, 387.0, 261.0, 172.0, 261.0, 267.0, 273.0, 184.0
-        ]
+        ])
     );
 }
 
@@ -79,10 +81,10 @@ fn conv_without_pad() {
     let session =
         pollster::block_on(wonnx::Session::from_model(conv_model)).expect("Session did not create");
     let result = pollster::block_on(session.run(&input_data)).unwrap();
-    assert_eq!(
-        result["Y"],
-        [54., 63., 72., 99., 108., 117., 144., 153., 162.]
-    );
+    common::assert_eq_vector(
+        (&result["Y"]).try_into().unwrap(),
+        &[54., 63., 72., 99., 108., 117., 144., 153., 162.],
+    )
 }
 
 #[test]
@@ -131,10 +133,12 @@ fn conv_stride() {
 
     let result = pollster::block_on(session.run(&input_data)).unwrap();
 
-    assert_eq!(
-        result["Y"],
-        [12., 27., 24., 63., 108., 81., 123., 198., 141., 112., 177., 124.]
-    );
+    common::assert_eq_vector(
+        (&result["Y"]).try_into().unwrap(),
+        &[
+            12., 27., 24., 63., 108., 81., 123., 198., 141., 112., 177., 124.,
+        ],
+    )
 }
 
 #[test]
@@ -173,7 +177,10 @@ fn conv_asymetric_stride() {
     let session =
         pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
     let result = pollster::block_on(session.run(&input_data)).unwrap();
-    assert_eq!(result["Y"], [21., 33., 99., 117., 189., 207., 171., 183.]);
+    common::assert_eq_vector(
+        (&result["Y"]).try_into().unwrap(),
+        &[21., 33., 99., 117., 189., 207., 171., 183.],
+    );
 }
 
 fn _conv_kernel_3() {
diff --git a/wonnx/tests/gather.rs b/wonnx/tests/gather.rs
index 00eb9281..a3cf7d00 100644
--- a/wonnx/tests/gather.rs
+++ b/wonnx/tests/gather.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, convert::TryInto};
 use wonnx::utils::{attribute, graph, model, node, tensor};
 mod common;
 
@@ -35,7 +35,7 @@ fn assert_gather(
         pollster::block_on(wonnx::Session::from_model(bn_model)).expect("Session did not create");
 
     let result = pollster::block_on(session.run(&input_data)).unwrap();
-    common::assert_eq_vector(result["Y"].as_slice(), output);
+    common::assert_eq_vector((&result["Y"]).try_into().unwrap(), output);
 }
 
 #[test]
diff --git a/wonnx/tests/globalaveragepool.rs b/wonnx/tests/globalaveragepool.rs
index 49dc63f8..1b7fb0ba 100644
--- a/wonnx/tests/globalaveragepool.rs
+++ b/wonnx/tests/globalaveragepool.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, convert::TryInto};
 use wonnx::utils::{graph, model, node, tensor};
 mod common;
 
@@ -49,5 +49,5 @@ fn global_average_pool() {
     // Channel 2: [[4,5], [6,7]] => average is 5,5
     // Channel 3: [[8,9], [10, 11]] => average is 9,5
     // Channel 4: [[12,13], [14, 15]] => average is 13,5
-    common::assert_eq_vector(out_y.as_slice(), &[1.5, 5.5, 9.5, 13.5]);
+    common::assert_eq_vector(out_y.try_into().unwrap(), &[1.5, 5.5, 9.5, 13.5]);
 }
diff --git a/wonnx/tests/identity.rs b/wonnx/tests/identity.rs
index a8a9f627..c9a671a0 100644
--- a/wonnx/tests/identity.rs
+++ b/wonnx/tests/identity.rs
@@ -1,5 +1,6 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, convert::TryInto};
 use wonnx::utils::{graph, model, node, tensor};
+mod common;
 
 #[test]
 fn test_identity() {
@@ -23,7 +24,7 @@ fn test_identity() {
         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"], data);
+    common::assert_eq_vector((&result["Y"]).try_into().unwrap(), &data);
 }
 
 #[test]
@@ -51,5 +52,5 @@ fn test_double_identity() {
         pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
 
     let result = pollster::block_on(session.run(&input_data)).unwrap();
-    assert_eq!(result["Z"], data);
+    common::assert_eq_vector((&result["Z"]).try_into().unwrap(), &data);
 }
diff --git a/wonnx/tests/matrix.rs b/wonnx/tests/matrix.rs
index 9196600a..eceded10 100644
--- a/wonnx/tests/matrix.rs
+++ b/wonnx/tests/matrix.rs
@@ -1,5 +1,6 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, convert::TryInto};
 use wonnx::utils::{attribute, graph, initializer, model, node, tensor};
+mod common;
 
 #[test]
 fn test_matmul_square_matrix() {
@@ -34,10 +35,9 @@ fn test_matmul_square_matrix() {
 
     let session =
         pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
-
     let result = pollster::block_on(session.run(&input_data)).unwrap();
 
-    assert_eq!(result["C"].as_slice(), sum.as_slice().unwrap());
+    common::assert_eq_vector((&result["C"]).try_into().unwrap(), sum.as_slice().unwrap());
 }
 
 #[test]
@@ -74,7 +74,7 @@ fn test_two_transposes() {
         pollster::block_on(wonnx::Session::from_model(model)).expect("session did not create");
     let result = pollster::block_on(session.run(&input_data)).unwrap();
 
-    assert_eq!(result["Z"], data);
+    common::assert_eq_vector((&result["Z"]).try_into().unwrap(), &data);
 }
 
 #[test]
@@ -102,9 +102,9 @@ fn test_split() {
     let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     let test_y = vec![1., 2., 3., 7., 8., 9.];
-    assert_eq!(result["Y"], test_y);
+    common::assert_eq_vector((&result["Y"]).try_into().unwrap(), &test_y);
     let test_w = vec![4., 5., 6., 10., 11., 12.];
-    assert_eq!(result["W"], test_w);
+    common::assert_eq_vector((&result["W"]).try_into().unwrap(), &test_w);
 }
 
 #[test]
@@ -133,7 +133,7 @@ fn test_resize() {
     let result = pollster::block_on(session.run(&input_data)).unwrap();
 
     let test_y = vec![1., 3.];
-    assert_eq!(result["Y"], test_y);
+    common::assert_eq_vector((&result["Y"]).try_into().unwrap(), &test_y);
 
     let mut input_data = HashMap::new();
     let data = (1..=4).map(|x| x as f32).collect::<Vec<f32>>();
diff --git a/wonnx/tests/onehot.rs b/wonnx/tests/onehot.rs
index fe0b805d..0ca20702 100644
--- a/wonnx/tests/onehot.rs
+++ b/wonnx/tests/onehot.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, convert::TryInto};
 use wonnx::{
     onnx::AttributeProto,
     utils::{attribute, graph, model, node, tensor, InputTensor},
@@ -50,7 +50,7 @@ fn test_onehot(
 
     let result = pollster::block_on(session.run(&input_data)).unwrap();
     log::info!("OUT: {:?}", result["Y"]);
-    common::assert_eq_vector(result["Y"].as_slice(), output);
+    common::assert_eq_vector((&result["Y"]).try_into().unwrap(), output);
 }
 
 #[test]
diff --git a/wonnx/tests/pretrained_models.rs b/wonnx/tests/pretrained_models.rs
index ada8f210..523ea4a9 100644
--- a/wonnx/tests/pretrained_models.rs
+++ b/wonnx/tests/pretrained_models.rs
@@ -1,7 +1,10 @@
 use image::{imageops::FilterType, ImageBuffer, Pixel, Rgb};
 use ndarray::s;
 use std::collections::HashMap;
+use std::convert::TryInto;
 use std::path::Path;
+use wonnx::utils::InputTensor;
+mod common;
 
 #[test]
 fn test_relu() {
@@ -13,19 +16,18 @@ fn test_relu() {
         .expect("session did not create");
     let result = pollster::block_on(session.run(&input_data)).unwrap();
 
-    assert_eq!(result["y"], &[0.0, 1.0]);
+    common::assert_eq_vector((&result["y"]).try_into().unwrap(), &[0.0, 1.0]);
 }
 
-#[test]
-fn test_mnist() {
-    let _ = env_logger::builder().is_test(true).try_init();
-    let image = load_image("0.jpg");
-    let mut input_data = HashMap::new();
-    input_data.insert("Input3".to_string(), image.as_slice().unwrap().into());
+fn infer_mnist(image: InputTensor) -> (usize, f32) {
     let session = pollster::block_on(wonnx::Session::from_path("../data/models/opt-mnist.onnx"))
         .expect("Session did not create");
 
-    let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
+    let mut input_data = HashMap::new();
+    input_data.insert("Input3".to_string(), image);
+    let output = pollster::block_on(session.run(&input_data)).unwrap();
+    let output_tensor: &[f32] = (&output["Plus214_Output_0"]).try_into().unwrap();
+    output_tensor
         .iter()
         .enumerate()
         .fold((0, 0.), |(idx_max, val_max), (idx, val)| {
@@ -34,56 +36,26 @@ fn test_mnist() {
             } else {
                 (idx, *val)
             }
-        });
+        })
+}
 
+#[test]
+fn test_mnist() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let image = load_image("0.jpg");
+    let result = infer_mnist(image.as_slice().unwrap().into());
     assert_eq!(result.0, 0);
 
     let image = load_image("3.jpg");
-    let mut input_data = HashMap::new();
-    input_data.insert("Input3".to_string(), image.as_slice().unwrap().into());
-    let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
-        .iter()
-        .enumerate()
-        .fold((0, 0.), |(idx_max, val_max), (idx, val)| {
-            if &val_max > val {
-                (idx_max, val_max)
-            } else {
-                (idx, *val)
-            }
-        });
-
+    let result = infer_mnist(image.as_slice().unwrap().into());
     assert_eq!(result.0, 3);
 
     let image = load_image("5.jpg");
-    let mut input_data = HashMap::new();
-    input_data.insert("Input3".to_string(), image.as_slice().unwrap().into());
-    let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
-        .iter()
-        .enumerate()
-        .fold((0, 0.), |(idx_max, val_max), (idx, val)| {
-            if &val_max > val {
-                (idx_max, val_max)
-            } else {
-                (idx, *val)
-            }
-        });
-
+    let result = infer_mnist(image.as_slice().unwrap().into());
     assert_eq!(result.0, 5);
 
     let image = load_image("7.jpg");
-    let mut input_data = HashMap::new();
-    input_data.insert("Input3".to_string(), image.as_slice().unwrap().into());
-    let result = pollster::block_on(session.run(&input_data)).unwrap()["Plus214_Output_0"]
-        .iter()
-        .enumerate()
-        .fold((0, 0.), |(idx_max, val_max), (idx, val)| {
-            if &val_max > val {
-                (idx_max, val_max)
-            } else {
-                (idx, *val)
-            }
-        });
-
+    let result = infer_mnist(image.as_slice().unwrap().into());
     assert_eq!(result.0, 7);
 }
 
@@ -100,9 +72,10 @@ fn test_squeeze() {
         .expect("session did not create");
     let result =
         &pollster::block_on(session.run(&input_data)).unwrap()["squeezenet0_flatten0_reshape0"];
-    let mut probabilities = result.iter().enumerate().collect::<Vec<_>>();
+    let output_tensor: &[f32] = result.try_into().unwrap();
+    let mut probabilities: Vec<(usize, f32)> = output_tensor.iter().cloned().enumerate().collect();
 
-    probabilities.sort_unstable_by(|a, b| b.1.partial_cmp(a.1).unwrap());
+    probabilities.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
 
     assert_eq!(probabilities[0].0, 22);
 }
diff --git a/wonnx/tests/reduce.rs b/wonnx/tests/reduce.rs
index 6e40f9e7..5391e196 100644
--- a/wonnx/tests/reduce.rs
+++ b/wonnx/tests/reduce.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, convert::TryInto};
 use wonnx::{
     onnx::AttributeProto,
     utils::{attribute, graph, initializer_int64, model, node, tensor},
@@ -38,7 +38,7 @@ fn test_reduce(
 
     let result = pollster::block_on(session.run(&input_data)).unwrap();
     log::info!("OUT: {:?}", result["Y"]);
-    common::assert_eq_vector(result["Y"].as_slice(), output);
+    common::assert_eq_vector((&result["Y"]).try_into().unwrap(), output);
 }
 
 fn sum_square(a: f32, b: f32) -> f32 {
@@ -276,5 +276,8 @@ fn test_reduce_sum_with_axes_as_input() {
 
     let result = pollster::block_on(session.run(&input_data)).unwrap();
     log::info!("OUT: {:?}", result["Y"]);
-    common::assert_eq_vector(result["Y"].as_slice(), &[4., 6., 12., 14., 20., 22.]);
+    common::assert_eq_vector(
+        (&result["Y"]).try_into().unwrap(),
+        &[4., 6., 12., 14., 20., 22.],
+    );
 }