Builder API and further removal of ONNX from the IR

wonnx-cli/src/cpu.rs

@@ -5,7 +5,7 @@ use async_trait::async_trait;
 use tract_onnx::prelude::*;
 use wonnx::{
     onnx::ModelProto,
-    utils::{OutputTensor, Shape},
+    tensor::{Shape, TensorData},
 };
 
 type RunnableOnnxModel =
@@ -73,7 +73,7 @@ impl Inferer for CPUInferer {
         outputs: &[String],
         inputs: &HashMap<String, crate::Tensor>,
         model: &ModelProto,
-    ) -> Result<HashMap<String, OutputTensor>, NNXError> {
+    ) -> Result<HashMap<String, TensorData<'static>>, NNXError> {
         let mut cpu_inputs: HashMap<usize, tract_onnx::prelude::Tensor> = HashMap::new();
 
         for (input_name, input_tensor) in inputs {
@@ -86,12 +86,7 @@ impl Inferer for CPUInferer {
                 .unwrap_or_else(|| panic!("input not found with name {}", input_name));
             log::info!("set input fact {} for cpu model", input_index.0,);
 
-            let dims: Vec<usize> = self.input_shapes[input_name]
-                .dims
-                .iter()
-                .map(|x| (*x) as usize)
-                .collect();
-
+            let dims: Vec<usize> = self.input_shapes[input_name].dims.to_vec();
             cpu_inputs.insert(input_index.0, input_tensor.to_tract_tensor(&dims)?);
         }
 
@@ -103,7 +98,7 @@ impl Inferer for CPUInferer {
         let result = self.model.run(cpu_inputs_ordered)?;
         log::debug!("cpu result: {:?}", result);
 
-        let mut output_tensors = HashMap::<String, OutputTensor>::new();
+        let mut output_tensors = HashMap::<String, TensorData>::new();
 
         for output_name in outputs {
             let result_vector = {
@@ -132,7 +127,7 @@ impl Inferer for CPUInferer {
             let av = result_vector.to_array_view()?;
             output_tensors.insert(
                 output_name.clone(),
-                OutputTensor::F32(av.as_slice().unwrap().to_vec()),
+                TensorData::F32(av.as_slice().unwrap().into()).into_static(),
             );
         }
         Ok(output_tensors)

wonnx-cli/src/gpu.rs

@@ -4,7 +4,7 @@ use wonnx::onnx::ModelProto;
 use wonnx::SessionConfig;
 
 use async_trait::async_trait;
-use wonnx::utils::OutputTensor;
+use wonnx::tensor::TensorData;
 
 use crate::types::Inferer;
 use crate::types::NNXError;
@@ -33,14 +33,14 @@ impl Inferer for GPUInferer {
         outputs: &[String],
         inputs: &HashMap<String, crate::Tensor>,
         _model: &ModelProto,
-    ) -> Result<HashMap<String, OutputTensor>, NNXError> {
+    ) -> Result<HashMap<String, TensorData<'static>>, NNXError> {
         let input_refs = inputs
             .iter()
             .map(|(k, v)| (k.clone(), v.input_tensor()))
             .collect();
         let mut result = self.session.run(&input_refs).await.expect("run failed");
 
-        let mut output_tensors = HashMap::<String, OutputTensor>::new();
+        let mut output_tensors = HashMap::<String, TensorData>::new();
 
         for output_name in outputs {
             let result = match result.remove(output_name) {

wonnx-cli/src/info.rs

@@ -4,7 +4,7 @@ use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use prettytable::{row, table, Table};
 use wonnx::{
     onnx::{GraphProto, ModelProto, NodeProto, ValueInfoProto},
-    utils::{ScalarType, Shape},
+    tensor::{ScalarType, Shape},
     WonnxError,
 };
 
@@ -31,7 +31,14 @@ fn dimensions_infos(
     }
 
     for info in graph_proto.get_initializer() {
-        let shape = Shape::from(ScalarType::from_i32(info.get_data_type())?, info.get_dims());
+        let shape = Shape::from(
+            ScalarType::from_onnx_i32(info.get_data_type())?,
+            &info
+                .get_dims()
+                .iter()
+                .map(|x| *x as usize)
+                .collect::<Vec<usize>>(),
+        );
         shapes_info.insert(info.get_name().to_string(), Some(shape));
     }
 
@@ -112,7 +119,14 @@ pub fn sizes_table(model: &ModelProto) -> Result<Table, WonnxError> {
 
     let mut initializer_size: usize = 0;
     for info in model.get_graph().get_initializer() {
-        let shape = Shape::from(ScalarType::from_i32(info.get_data_type())?, info.get_dims());
+        let shape = Shape::from(
+            ScalarType::from_onnx_i32(info.get_data_type())?,
+            &info
+                .get_dims()
+                .iter()
+                .map(|x| *x as usize)
+                .collect::<Vec<usize>>(),
+        );
         initializer_size += shape.buffer_bytes_aligned();
     }
 

wonnx-cli/src/main.rs

@@ -8,7 +8,8 @@ use std::fs::File;
 use structopt::StructOpt;
 use trace::trace_command;
 use wonnx::onnx::ModelProto;
-use wonnx::utils::{get_opset_version, OutputTensor, Shape};
+use wonnx::onnx_model::get_opset_version;
+use wonnx::tensor::{Shape, TensorData};
 use wonnx_preprocessing::shape_inference::{apply_dynamic_dimensions, infer_shapes};
 use wonnx_preprocessing::text::{get_lines, EncodedText};
 use wonnx_preprocessing::Tensor;
@@ -105,7 +106,7 @@ async fn run() -> Result<(), NNXError> {
 fn print_qa_output(
     infer_opt: &InferOptions,
     qa_encoding: &EncodedText,
-    mut outputs: HashMap<String, OutputTensor>,
+    mut outputs: HashMap<String, TensorData>,
 ) -> Result<(), NNXError> {
     let start_output: Vec<f32> = outputs
         .remove(&infer_opt.qa_answer_start)
@@ -133,7 +134,7 @@ fn print_qa_output(
 fn print_output(
     infer_opt: &InferOptions,
     output_name: &str,
-    output: OutputTensor,
+    output: TensorData,
     print_output_names: bool,
     print_newlines: bool,
 ) {
@@ -165,35 +166,35 @@ fn print_output(
 
             // Just print the output tensor values, one a line
             match output {
-                wonnx::utils::OutputTensor::F32(fs) => {
-                    for i in fs {
+                wonnx::tensor::TensorData::F32(fs) => {
+                    for i in fs.iter() {
                         if print_newlines {
                             println!("{:.3}", i);
                         } else {
                             print!("{:.3} ", i);
                         }
                     }
                 }
-                wonnx::utils::OutputTensor::I32(ints) => {
-                    for i in ints {
+                wonnx::tensor::TensorData::I32(ints) => {
+                    for i in ints.iter() {
                         if print_newlines {
                             println!("{}", i);
                         } else {
                             print!("{}", i);
                         }
                     }
                 }
-                wonnx::utils::OutputTensor::I64(ints) => {
-                    for i in ints {
+                wonnx::tensor::TensorData::I64(ints) => {
+                    for i in ints.iter() {
                         if print_newlines {
                             println!("{}", i);
                         } else {
                             print!("{}", i);
                         }
                     }
                 }
-                wonnx::utils::OutputTensor::U8(ints) => {
-                    for i in ints {
+                wonnx::tensor::TensorData::U8(ints) => {
+                    for i in ints.iter() {
                         if print_newlines {
                             println!("{}", i);
                         } else {
@@ -238,7 +239,10 @@ async fn prepare_command(prepare_opt: PrepareOptions) -> Result<(), NNXError> {
             {
                 Some(input) => {
                     let data_type = input.data_type().map_err(|_| NNXError::InvalidInputShape)?;
-                    let shape = Shape::from(data_type, &new_dims);
+                    let shape = Shape::from(
+                        data_type,
+                        &new_dims.iter().map(|x| *x as usize).collect::<Vec<usize>>(),
+                    );
                     input.set_shape(&shape);
                     log::info!("setting shape of input {input_name} to {shape}");
                 }

wonnx-cli/src/types.rs

@@ -4,7 +4,8 @@ use structopt::StructOpt;
 use thiserror::Error;
 use wonnx::{
     onnx::ModelProto,
-    utils::{OpsetError, OutputTensor, Shape, TensorConversionError},
+    onnx_model::OpsetError,
+    tensor::{Shape, TensorConversionError, TensorData},
     SessionError, WonnxError,
 };
 use wonnx_preprocessing::{
@@ -296,7 +297,7 @@ pub trait Inferer {
         outputs: &[String],
         inputs: &HashMap<String, Tensor>,
         model: &ModelProto,
-    ) -> Result<HashMap<String, OutputTensor>, NNXError>;
+    ) -> Result<HashMap<String, TensorData<'static>>, NNXError>;
 }
 
 pub struct InferenceInput {

wonnx-cli/src/utils.rs

@@ -2,7 +2,7 @@ use ndarray::{Array, ArrayBase};
 use std::collections::HashMap;
 use std::path::Path;
 use wonnx::onnx::{ModelProto, TensorShapeProto, ValueInfoProto};
-use wonnx::utils::{DataTypeError, ScalarType, Shape};
+use wonnx::tensor::{DataTypeError, ScalarType, Shape};
 use wonnx::WonnxError;
 use wonnx_preprocessing::image::{load_bw_image, load_rgb_image};
 use wonnx_preprocessing::text::{EncodedText, TextTokenizer};
@@ -55,7 +55,7 @@ impl ValueInfoProtoUtil for ValueInfoProto {
         Ok(match &self.get_field_type().value {
             Some(x) => match x {
                 wonnx::onnx::TypeProto_oneof_value::tensor_type(t) => {
-                    ScalarType::from_i32(t.get_elem_type())?
+                    ScalarType::from_onnx_i32(t.get_elem_type())?
                 }
                 wonnx::onnx::TypeProto_oneof_value::sequence_type(_) => todo!(),
                 wonnx::onnx::TypeProto_oneof_value::map_type(_) => todo!(),
@@ -119,8 +119,8 @@ pub fn load_image_input(
     input_shape: &Shape,
 ) -> Result<ArrayBase<ndarray::OwnedRepr<f32>, ndarray::IxDyn>, NNXError> {
     if input_shape.rank() == 3 {
-        let mut w = input_shape.dim(1) as usize;
-        let mut h = input_shape.dim(2) as usize;
+        let mut w = input_shape.dim(1);
+        let mut h = input_shape.dim(2);
         if w == 0 {
             w = 224;
         }
@@ -138,8 +138,8 @@ pub fn load_image_input(
             Err(NNXError::InvalidInputShape)
         }
     } else if input_shape.rank() == 4 {
-        let mut w = input_shape.dim(2) as usize;
-        let mut h = input_shape.dim(3) as usize;
+        let mut w = input_shape.dim(2);
+        let mut h = input_shape.dim(3);
         if w == 0 {
             w = 224;
         }
@@ -179,12 +179,12 @@ impl InferenceInput {
                 .get_input_shape(&infer_opt.qa_segment_input)?
                 .ok_or_else(|| NNXError::InputNotFound(infer_opt.qa_segment_input.clone()))?;
 
-            let segment_length = tokens_input_shape.element_count() as usize;
+            let segment_length = tokens_input_shape.element_count();
 
-            if segment_length != mask_input_shape.element_count() as usize {
+            if segment_length != mask_input_shape.element_count() {
                 return Err(NNXError::InvalidInputShape);
             }
-            if segment_length != segment_input_shape.element_count() as usize {
+            if segment_length != segment_input_shape.element_count() {
                 return Err(NNXError::InvalidInputShape);
             }
 
@@ -272,7 +272,7 @@ impl InferenceInput {
 
             let values: Result<Vec<f32>, _> = text.split(',').map(|v| v.parse::<f32>()).collect();
             let mut values = values.map_err(NNXError::InvalidNumber)?;
-            values.resize(raw_input_shape.element_count() as usize, 0.0);
+            values.resize(raw_input_shape.element_count(), 0.0);
             inputs.insert(
                 raw_input_name.clone(),
                 Tensor::F32(Array::from_vec(values).into_dyn()),