Repository: webonnx/wonnx (language: python)

Issue:
feat: implement a very slow softmax
Currently the Softmax operator just copies data from input to output. To match Keras, I implemented actual Softmax in WGSL. 

This is the first time I have written any WGSL and my suspicion is that this is probably the slowest possible way to implement it (i.e. single thread, fully serial, not running anything in parallel on the GPU at all - main CPU is probably faster at this!). Nevertheless I think it is good to have this at this point (can make it fast later). 

Note, the actual ONNX spec may also require Softmax on a single axis in a matrix, see [here](https://github.com/onnx/onnx/blob/master/onnx/backend/test/case/node/softmax.py). Probably also should write some tests.
Replace panicking with proper error handling
This also has some of my other changes, it turned out quite impossible to do this separately from the other changes as it is quite invasive. I suggest you have a look at my other PRs first, then merge those, and rebase this on top of that before merging.
feat: implement a very slow softmax
Currently the Softmax operator just copies data from input to output. To match Keras, I implemented actual Softmax in WGSL. 

This is the first time I have written any WGSL and my suspicion is that this is probably the slowest possible way to implement it (i.e. single thread, fully serial, not running anything in parallel on the GPU at all - main CPU is probably faster at this!). Nevertheless I think it is good to have this at this point (can make it fast later). 

Note, the actual ONNX spec may also require Softmax on a single axis in a matrix, see [here](https://github.com/onnx/onnx/blob/master/onnx/backend/test/case/node/softmax.py). Probably also should write some tests.


Discussion / hints:

Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index ab85fec2..937764a2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1898,6 +1898,7 @@ dependencies = [
  "pollster",
  "protobuf",
  "tera",
+ "thiserror",
  "ureq",
  "wasm-bindgen",
  "wasm-bindgen-futures",
diff --git a/Cargo.toml b/Cargo.toml
index 72dfa27f..69119de6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -32,6 +32,7 @@ log = "0.4.14"
 protobuf = { version = "2", features = ["with-bytes"] }
 tera = "1.15.0"
 lazy_static = "1.4.0"
+thiserror = "^1.0.0"
 
 [dev-dependencies]
 image = "0.23.14"
diff --git a/examples/simple_graph.rs b/examples/simple_graph.rs
index 6a243d10..c9e8d3f0 100644
--- a/examples/simple_graph.rs
+++ b/examples/simple_graph.rs
@@ -1,19 +1,21 @@
 use std::collections::HashMap;
 // use wasm_bindgen_test::*;
 // Indicates a f32 overflow in an intermediate Collatz value
-use std::error;
-use wonnx::utils::{attribute, graph, initializer, model, node, tensor};
-type Result<T> = std::result::Result<T, Box<dyn error::Error>>;
+use wonnx::{
+    utils::{attribute, graph, initializer, model, node, tensor},
+    SessionError, WonnxError,
+};
 
-async fn run() {
-    let result = &execute_gpu().await.unwrap();
+async fn run() -> Result<(), WonnxError> {
+    let result = &execute_gpu().await?;
     let (_, result) = result.iter().next().unwrap();
 
     assert_eq!(result, &[54., 63., 72., 99., 108., 117., 144., 153., 162.]);
+    Ok(())
 }
 
 // Hardware management
-async fn execute_gpu() -> Result<HashMap<String, Vec<f32>>> {
+async fn execute_gpu() -> Result<HashMap<String, Vec<f32>>, SessionError> {
     // USER INPUT
     let n = 5;
     let c = 1;
@@ -55,7 +57,7 @@ fn main() {
     #[cfg(not(target_arch = "wasm32"))]
     {
         env_logger::init();
-        pollster::block_on(run());
+        pollster::block_on(run()).unwrap();
     }
     #[cfg(target_arch = "wasm32")]
     {
diff --git a/examples/squeeze.rs b/examples/squeeze.rs
index c5bd5c59..acfbf105 100644
--- a/examples/squeeze.rs
+++ b/examples/squeeze.rs
@@ -4,7 +4,6 @@ use std::collections::HashMap;
 use image::{imageops::FilterType, ImageBuffer, Pixel, Rgb};
 use log::info;
 use ndarray::s;
-use std::error;
 use std::time::Instant;
 use std::{
     fs,
@@ -12,8 +11,8 @@ use std::{
     path::Path,
     time::Duration,
 };
+use wonnx::WonnxError;
 
-type Result<T> = std::result::Result<T, Box<dyn error::Error>>;
 // Args Management
 async fn run() {
     let probabilities = &execute_gpu().await.unwrap();
@@ -39,23 +38,21 @@ async fn run() {
 }
 
 // Hardware management
-async fn execute_gpu() -> Result<HashMap<String, Vec<f32>>> {
+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());
 
-    let session = wonnx::Session::from_path("examples/data/models/opt-squeeze.onnx")
-        .await
-        .unwrap();
+    let session = wonnx::Session::from_path("examples/data/models/opt-squeeze.onnx").await?;
     let time_pre_compute = Instant::now();
     info!("Start Compute");
-    let result = session.run(input_data.clone()).await;
+    let result = session.run(input_data.clone()).await?;
     let time_post_compute = Instant::now();
     println!(
         "time: first_prediction: {:#?}",
         time_post_compute - time_pre_compute
     );
-    result
+    Ok(result)
 }
 
 fn main() {
@@ -88,7 +85,7 @@ pub fn load_image() -> ndarray::ArrayBase<ndarray::OwnedRepr<f32>, ndarray::Dim<
 
     let image_buffer: ImageBuffer<Rgb<u8>, Vec<u8>> = image::open(image_path)
         .unwrap()
-        .resize_to_fill(224 as u32, 224 as u32, FilterType::Nearest)
+        .resize_to_fill(224, 224, FilterType::Nearest)
         .to_rgb8();
 
     // Python:
diff --git a/src/compiler.rs b/src/compiler.rs
index 48f8f09e..5aa03d7c 100644
--- a/src/compiler.rs
+++ b/src/compiler.rs
@@ -1,6 +1,7 @@
-use crate::utils::{ceil, get_attribute};
+use crate::utils::{ceil, get_attribute, AttributeNotFoundError};
 use std::collections::HashMap;
 use tera::{Context, Tera};
+use thiserror::Error;
 
 // Escaping special characters as well as adding `var_` in the beginning of the variable name to avoid collisions with wgsl syntax.
 // FIXME: this has the potential to cause collisions (i.e. "a/b" and "a.b" will both translate to "ab" and hence cannot be used simultaneously)
@@ -14,31 +15,66 @@ pub struct CompiledNode {
     pub threads: (u32, u32, u32),
 }
 
+#[derive(Error, Debug)]
+pub enum CompileError {
+    #[error("dimensions information missing for input/output '{0}' of node '{1}'. You may want to run onnx-simplifier on the model first.")]
+    DimensionsMissing(String, String),
+
+    #[error("attribute not found: {0}")]
+    AttributeNotFound(#[from] AttributeNotFoundError),
+
+    #[error("operation not recognized: {0}")]
+    InvalidOperation(String),
+
+    #[error("op {0} is not implemented yet! Check the README if you want to implement it")]
+    UnimplementedOp(String),
+
+    #[error("'{variant}' is not yet implemented for op {op}")]
+    UnimplementedVariant { variant: String, op: String },
+
+    #[error("the opset version {0} is not supported")]
+    UnsupportedOpsetVersion(i64),
+
+    #[error("the value '{attribute}' is invalid for attribute '{value}' (opset version {opset_version})")]
+    InvalidAttributeValue {
+        attribute: String,
+        value: String,
+        opset_version: i64,
+    },
+}
+
 pub fn compile(
     node: &crate::onnx::NodeProto,
     dims_infos: &HashMap<String, Vec<i64>>,
     tera: &Tera,
-) -> CompiledNode {
+    opset_version: i64,
+) -> Result<CompiledNode, CompileError> {
     // Escape unwanted characters
     let mut inputs = node.get_input().to_vec();
     let mut outputs = node.get_output().to_vec();
 
     let input_dims = inputs
         .iter()
-        .map(|input| {
-            dims_infos
-                .get(input.as_str())
-                .unwrap_or_else(|| panic!("Dimensions information not found for input '{}' of node '{}'. You may want to run onnx-simplifier on the model first.", input, node.get_name()))
+        .map(|input| match dims_infos.get(input.as_str()) {
+            Some(info) => Ok(info),
+            None => Err(CompileError::DimensionsMissing(
+                input.to_string(),
+                node.get_name().to_string(),
+            )),
         })
-        .collect::<Vec<_>>();
+        .collect::<Result<Vec<_>, CompileError>>()?;
+
     let output_dims = outputs
         .iter()
-        .map(|output| {
-            dims_infos
-                .get(output.as_str())
-                .unwrap_or_else(|| panic!("Dimensions information not found for output '{}' of mode '{}'. You may want to run onnx-simplifier on the model first.", output, node.get_name()))
+        .map(|output| match dims_infos.get(output.as_str()) {
+            Some(info) => Ok(info),
+            None => Err(CompileError::DimensionsMissing(
+                output.to_string(),
+                node.get_name().to_string(),
+            )),
         })
-        .collect::<Vec<_>>();
+        .collect::<Result<Vec<_>, CompileError>>()?;
+
     let input_lengths = input_dims
         .iter()
         .map(|dims| dims.iter().product())
@@ -90,6 +126,7 @@ pub fn compile(
     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() {
         // Map simple function
@@ -101,16 +138,58 @@ pub fn compile(
             1,
         ),
         // Copy data
-        "Reshape" | "Dropout" | "Flatten" | "Squeeze" | "Softmax" => (
+        "Reshape" | "Dropout" | "Flatten" | "Squeeze" => (
             "endomorphism/copy.wgsl".to_string(),
             ceil(output_lengths[0], 16) as _,
             1,
             1,
         ),
+        "Softmax" => {
+            let default_axis = match opset_version {
+                1..=10 => 1,   // https://github.com/onnx/onnx/blob/master/docs/Changelog.md#softmax-1
+                11..=12 => 1, // https://github.com/onnx/onnx/blob/master/docs/Changelog.md#softmax-11
+                13..=15 => -1, // https://github.com/onnx/onnx/blob/master/docs/Changelog.md#softmax-13
+                _ => return Err(CompileError::UnsupportedOpsetVersion(opset_version)),
+            };
+
+            /* Describes the axis of the inputs when coerced to 2D; defaults to one because the 0th axis most likely
+            describes the batch_size. From version 13 onwards, counting backwards is also allowed. We only support the
+            variant with [1,n] input tensors, where axis is 1 or -1 */
+            let mut axis = get_attribute("axis", Some(default_axis), node)?;
+            if axis < 0 {
+                if opset_version >= 13 {
+                    axis += input_dims.len() as i64;
+                } else {
+                    return Err(CompileError::InvalidAttributeValue {
+                        attribute: "axis".to_string(),
+                        value: format!("{}", axis),
+                        opset_version,
+                    });
+                }
+            }
+
+            if axis >= (input_dims.len() as i64) {
+                return Err(CompileError::InvalidAttributeValue {
+                    attribute: "axis".to_string(),
+                    value: format!("{}", axis),
+                    opset_version,
+                });
+            }
+
+            if axis != 1 {
+                return Err(CompileError::UnimplementedVariant {
+                    variant: "softmax on an axis other than the second with [1,n] inputs"
+                        .to_string(),
+                    op: "Softmax".to_string(),
+                });
+            }
+
+            ("endomorphism/softmax.wgsl".to_string(), 1, 1, 1)
+        }
         // Arithmetic operation
         "Add" | "And" | "Div" | "Equal" | "Greater" | "GreaterOrEqual" | "Less" | "LessOrEqual"
         | "Mod" | "Mul" | "Or" | "Sub" => {
-            let coefficient = get_attribute("coefficient", Some(1.0), node);
+            let coefficient = get_attribute("coefficient", Some(1.0), node)?;
             context.insert("coefficient", &coefficient);
             context.insert(
                 "op_type",
@@ -127,7 +206,11 @@ pub fn compile(
                     "Mul" => "*",
                     "Or" => "|",
                     "Sub" => "-",
-                    _ => unimplemented!(),
+                    _ => {
+                        return Err(CompileError::UnimplementedOp(
+                            node.get_op_type().to_string(),
+                        ))
+                    }
                 },
             );
             (
@@ -139,10 +222,12 @@ pub fn compile(
         }
         // Not taking into account attributes
         "BatchNormalization" => {
-            let epsilon = get_attribute("epsilon", Some(1.0), node);
+            let epsilon = get_attribute("epsilon", Some(1.0), node)?;
             context.insert("epsilon", &epsilon);
 
-            todo!();
+            return Err(CompileError::UnimplementedOp(
+                node.get_op_type().to_string(),
+            ));
 
             //   (
             //       "endomorphism/batchnormalization.wgsl".to_string(),
@@ -152,7 +237,7 @@ pub fn compile(
             //   )
         }
         "Relu" | "Sigmoid" | "Softsign" | "Softplus" | "Clip" | "Celu" | "Elu" | "LeakyRelu" => {
-            let alpha = get_attribute("alpha", Some(1.0), node);
+            let alpha = get_attribute("alpha", Some(1.0), node)?;
             context.insert("alpha", &alpha);
             (
                 "endomorphism/activation.wgsl".to_string(),
@@ -180,11 +265,11 @@ pub fn compile(
             // TODO: Conv only support NxCxHxW for the moment.
             debug_assert!(input_dims[0].len() == 4usize);
 
-            let auto_pad = get_attribute("auto_pad", Some("NOTSET".to_string()), node);
-            let dilations = get_attribute("dilations", Some(vec![1, 1]), node);
-            let kernel_shape = get_attribute::<Vec<i64>>("kernel_shape", None, node);
-            let strides = get_attribute("strides", Some(vec![1, 1]), node);
-            let pads = get_attribute("pads", Some(vec![0, 0, 0, 0]), node);
+            let auto_pad = get_attribute("auto_pad", Some("NOTSET".to_string()), node)?;
+            let dilations = get_attribute("dilations", Some(vec![1, 1]), node)?;
+            let kernel_shape = get_attribute::<Vec<i64>>("kernel_shape", None, node)?;
+            let strides = get_attribute("strides", Some(vec![1, 1]), node)?;
+            let pads = get_attribute("pads", Some(vec![0, 0, 0, 0]), node)?;
 
             let pads = match auto_pad.as_str() {
                 "NOTSET" => pads.to_vec(),
@@ -216,7 +301,12 @@ pub fn compile(
                         slack_1_div_2,
                     ]
                 }
-                _ => unimplemented!(),
+                _ => {
+                    return Err(CompileError::UnimplementedVariant {
+                        op: op.to_string(),
+                        variant: format!("auto_pad={}", auto_pad),
+                    })
+                }
             };
 
             let input_dims = input_dims[0];
@@ -246,7 +336,7 @@ pub fn compile(
                 ),
                 "Conv" | "ConvRelu" | "ConvLeakyRelu" | "ConvMish" => {
                     // Alpha is the Leaky Relu attribute
-                    let alpha = get_attribute("alpha", Some(0.01), node);
+                    let alpha = get_attribute("alpha", Some(0.01), node)?;
                     context.insert("alpha", &alpha);
 
                     // GLSL shader for convolution computation
@@ -282,12 +372,12 @@ pub fn compile(
                         )
                     }
                 }
-                _ => panic!("Invalid Opset"),
+                _ => return Err(CompileError::InvalidOperation(op.to_string())),
             }
         }
         "Gemm" | "MatMul" => {
-            let alpha = get_attribute("alpha", Some(1.0), node);
-            let beta = get_attribute("beta", Some(1.0), node);
+            let alpha = get_attribute("alpha", Some(1.0), node)?;
+            let beta = get_attribute("beta", Some(1.0), node)?;
             context.insert("alpha", &alpha);
             context.insert("beta", &beta);
 
@@ -304,7 +394,7 @@ pub fn compile(
                 "coordinate_transformation_mode",
                 Some("half_pixel".to_string()),
                 node,
-            );
+            )?;
             context.insert(
                 "coordinate_transformation_mode",
                 &coordinate_transformation_mode,
@@ -316,19 +406,26 @@ pub fn compile(
                 "align_corners" => {}
                 "asymmetric" => {}
                 "tf_crop_and_resize" => {
-                    let roi = get_attribute::<Vec<i64>>("roi", None, node);
-                    let extrapolation_value = get_attribute("extrapolation_value", Some(0.0), node);
+                    let roi = get_attribute::<Vec<i64>>("roi", None, node)?;
+                    let extrapolation_value =
+                        get_attribute("extrapolation_value", Some(0.0), node)?;
                     context.insert("roi", &roi);
                     context.insert("extrapolation_value", &extrapolation_value);
                 }
                 _ => {
-                    unimplemented!("This resize coordinate transformation is not implemented.")
+                    return Err(CompileError::UnimplementedVariant {
+                        op: "Resize".to_string(),
+                        variant: format!(
+                            "coordinate_transformation_mode={}",
+                            coordinate_transformation_mode
+                        ),
+                    })
                 }
             }
 
-            let scales = get_attribute::<Vec<f32>>("scales", Some(vec![]), node);
+            let scales = get_attribute::<Vec<f32>>("scales", Some(vec![]), node)?;
             let scale_prints = if scales.is_empty() {
-                let sizes = get_attribute::<Vec<i64>>("sizes", Some(vec![]), node);
+                let sizes = get_attribute::<Vec<i64>>("sizes", Some(vec![]), node)?;
                 sizes
                     .iter()
                     .enumerate()
@@ -341,33 +438,45 @@ pub fn compile(
                 scales.iter().map(|x| format!("{:.2}", x)).collect()
             };
 
-            let mode = get_attribute("mode", Some("nearest".to_string()), node);
+            let mode = get_attribute("mode", Some("nearest".to_string()), node)?;
             context.insert("mode", &mode);
             context.insert("scales", &scale_prints);
 
             match mode.as_str() {
                 "nearest" => {
-                    let nearest_mode =
-                        get_attribute("nearest_mode", Some("round_prefer_floor".to_string()), node);
+                    let nearest_mode = get_attribute(
+                        "nearest_mode",
+                        Some("round_prefer_floor".to_string()),
+                        node,
+                    )?;
                     match nearest_mode.as_str() {
                         "floor" => {}
-                        _ => unimplemented!(),
+                        _ => {
+                            return Err(CompileError::UnimplementedVariant {
+                                op: "Resize".to_string(),
+                                variant: format!("nearest_mode={}", nearest_mode.to_string()),
+                            })
+                        }
                     }
                 }
-                "linear" => {
-                    unimplemented!("Is not implemented yet");
-                }
                 "cubic" => {
-                    let cubic_coeff_a = get_attribute("cubic_coeff_a", Some(-0.75), node);
+                    let cubic_coeff_a = get_attribute("cubic_coeff_a", Some(-0.75), node)?;
                     context.insert("cubic_coeff_a", &cubic_coeff_a);
-                    unimplemented!("Is not implemented yet");
+                    return Err(CompileError::UnimplementedVariant {
+                        op: String::from("Resize"),
+                        variant: format!("mode={}", mode),
+                    });
                 }
+                /* "linear" | */
                 _ => {
-                    unimplemented!("This resize mode is not implemented.")
+                    return Err(CompileError::UnimplementedVariant {
+                        op: String::from("Resize"),
+                        variant: format!("mode={}", mode),
+                    });
                 }
             };
 
-            let exclude_outside = get_attribute("exclude_outside", Some(0), node);
+            let exclude_outside = get_attribute("exclude_outside", Some(0), node)?;
             context.insert("exclude_outside", &exclude_outside);
 
             (
@@ -377,11 +486,9 @@ pub fn compile(
                 1,
             )
         }
-        "Sum" => {
-            unimplemented!()
-        }
+        "Sum" => return Err(CompileError::UnimplementedOp(String::from("Sum"))),
         "Split" => {
-            let mut axis = get_attribute("axis", Some(0), node);
+            let mut axis = get_attribute("axis", Some(0), node)?;
             if axis < 0 {
                 axis += input_dims[0].len() as i64
             }
@@ -392,7 +499,7 @@ pub fn compile(
                 .map(|x| (x * split_chunk) as _)
                 .collect();
 
-            let split = get_attribute::<Vec<i64>>("split", Some(default_split), node);
+            let split = get_attribute::<Vec<i64>>("split", Some(default_split), node)?;
             context.insert("split", &split);
 
             (
@@ -404,7 +511,7 @@ pub fn compile(
         }
         "Transpose" => {
             let default = (input_lengths[0]..0).collect::<Vec<_>>();
-            let perms = get_attribute("perm", Some(default), node);
+            let perms = get_attribute("perm", Some(default), node)?;
             let permuted_dims = perms
                 .iter()
                 .map(|p| output_dims[0][*p as usize])
@@ -425,10 +532,7 @@ pub fn compile(
                 1,
             )
         }
-        op => panic!(
-            "Op {} is not implemented yet! Check the README if you want to implement it",
-            op
-        ),
+        op => return Err(CompileError::UnimplementedOp(op.to_string())),
     };
 
     let shader = tera
@@ -441,8 +545,8 @@ pub fn compile(
         node.get_name(),
         x - 16352
     );
-    CompiledNode {
+    Ok(CompiledNode {
         shader,
         threads: (x, y, z),
-    }
+    })
 }
diff --git a/src/lib.rs b/src/lib.rs
index 42b64bcb..bcae9fba 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,3 @@
-use std::error;
 pub mod compiler;
 pub mod onnx;
 pub mod optimisation;
@@ -9,13 +8,25 @@ pub mod utils;
 #[macro_use]
 extern crate lazy_static;
 
-use optimisation::EncoderBuilder;
-use protobuf::{self, Message};
+use compiler::CompileError;
+use onnx::ValueInfoProto;
+use optimisation::{EncoderBuilder, OptimizationError};
+use protobuf::{self, Message, ProtobufError};
 use std::collections::HashMap;
+use std::result::Result;
+use utils::len;
 
 use crate::resource::resize;
-// Change the alias to `Box<error::Error>`.
-type Result<T> = std::result::Result<T, Box<dyn error::Error>>;
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum WonnxError {
+    #[error("error compiling model: {0}")]
+    CompileError(#[from] CompileError),
+
+    #[error("error executing the model: {0}")]
+    SessionError(#[from] SessionError),
+}
 
 /// Creates a new session connected to the GPU.
 ///
@@ -52,36 +63,77 @@ type Result<T> = std::result::Result<T, Box<dyn error::Error>>;
 pub struct Session {
     pub device: wgpu::Device,
     pub queue: wgpu::Queue,
-    pub inner_infos: HashMap<String, wgpu::Buffer>,
+    pub buffers: HashMap<String, wgpu::Buffer>,
     pub builders: Vec<EncoderBuilder>,
-    pub outputs: Vec<String>,
+    pub outputs: Vec<ValueInfoProto>,
+}
+
+#[derive(Error, Debug)]
+pub enum SessionError {
+    #[error("could not deserialize model: {0}")]
+    ModelDeserializationError(#[from] ProtobufError),
+
+    #[error("an error occurred reading the model file: {0}")]
+    ModelReadingError(#[from] std::io::Error),
+
+    #[error(
+        "invalid input name '{0}'; inspect the file with e.g. Netron to find the correct name"
+    )]
+    InvalidInput(String),
+
+    #[error(
+        "invalid output name '{0}'; inspect the file with e.g. Netron to find the correct name"
+    )]
+    InvalidOutput(String),
+
+    #[error("optimization failed: {0}")]
+    OptimizationFailed(#[from] OptimizationError),
 }
 
 impl Session {
     // Read an ONNX model from a path and create a session.
-    pub async fn from_path(path: &str) -> Result<Session> {
-        let model = onnx::ModelProto::parse_from_bytes(
-            &std::fs::read(path).expect("ONNX Model path not found."),
-        )
-        .expect("Could not deserialize the Model");
+    pub async fn from_path(path: &str) -> Result<Session, SessionError> {
+        let model = onnx::ModelProto::parse_from_bytes(&std::fs::read(path)?)?;
 
         Session::from_model(model).await
     }
 
     // Create a Session given an ONNX model.
-    pub async fn from_model(model: onnx::ModelProto) -> Result<Session> {
+    pub async fn from_model(model: onnx::ModelProto) -> Result<Session, SessionError> {
         let promise = resource::request_device_queue();
-
         let (device, queue) = promise.await;
 
-        let (inner_infos, builders) = optimisation::load(model.get_graph(), &device).unwrap();
+        // Find the version of the ONNX operator set this model is using (this is useful because some operators' specifications change over time).
+        // Note, if any other op set than the ONNX operator set is referenced, we cannot run the model.
+        // See https://github.com/onnx/onnx/blob/master/docs/Versioning.md#operator-sets
+        let mut onnx_opset_version = None;
+        for opset_import in model.get_opset_import() {
+            match opset_import.get_domain() {
+                "" => {
+                    // This is a reference to the ONNX specification op set
+                    if let Some(onnx_version) = onnx_opset_version {
+                        panic!(
+                            "two onnx operator set versions specified (first is {}, second is {})",
+                            onnx_version,
+                            opset_import.get_version()
+                        );
+                    } else {
+                        onnx_opset_version = Some(opset_import.get_version());
+                    }
+                }
+                some_other_opset => {
+                    panic!(
+                        "the model references other op sets than the ONNX specified opset ({})",
+                        some_other_opset
+                    );
+                }
+            }
+        }
 
+        let onnx_opset_version =
+            onnx_opset_version.expect("no ONNX opset was referenced by the model!");
         let graph = model.get_graph();
-        let outputs = graph
-            .get_output()
-            .iter()
-            .map(|info| (info.get_name().to_string()))
-            .collect();
+        let (buffers, builders) = optimisation::load(graph, &device, onnx_opset_version).unwrap();
 
         // The data is loaded after the first submit
         let encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
@@ -90,9 +142,9 @@ impl Session {
         Ok(Session {
             device,
             queue,
-            inner_infos,
+            buffers,
             builders,
-            outputs,
+            outputs: graph.get_output().to_vec(),
         })
     }
 
@@ -100,22 +152,22 @@ impl Session {
     // It copy input data to the buffers.
     // Run the command encoder.
     // Copy the output into an exit buffer that can be deleted.
-    pub async fn run(&self, inputs: HashMap<String, &[f32]>) -> Result<HashMap<String, Vec<f32>>> {
+    pub async fn run(
+        &self,
+        inputs: HashMap<String, &[f32]>,
+    ) -> Result<HashMap<String, Vec<f32>>, SessionError> {
         let device = &self.device;
         let queue = &self.queue;
         let builders = &self.builders;
-        let inner_infos = &self.inner_infos;
+        let inner_infos = &self.buffers;
         let outputs = &self.outputs;
 
         // Copy input data
         for (input, data) in inputs {
             queue.write_buffer(
-                inner_infos.get(&input).unwrap_or_else(|| {
-                    panic!(
-                        "Invalid input: {}, try to use netron.app to see the correct input name",
-                        input
-                    )
-                }),
+                inner_infos
+                    .get(&input)
+                    .ok_or(SessionError::InvalidInput(input))?,
                 0,
                 bytemuck::cast_slice(&resize(data.to_vec())),
             )
@@ -142,8 +194,12 @@ impl Session {
         let mut results = HashMap::new();
 
         for output in outputs {
+            let output_name = output.get_name();
+
             // Copy the output data into the staging buffer.
-            let buffer = inner_infos.get(output.as_str()).unwrap();
+            let buffer = inner_infos
+                .get(output_name)
+                .ok_or_else(|| SessionError::InvalidOutput(output_name.to_string()))?;
 
             let buffer_slice = buffer.slice(..);
             // TODO: Define behavior for multi output.
@@ -155,7 +211,15 @@ impl Session {
             // Gets contents of buffer
             let data = buffer_slice.get_mapped_range();
             // Since contents are got in bytes, this converts these bytes back to f32
-            results.insert(output.clone(), bytemuck::cast_slice(&data).to_vec());
+
+            // 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 = len(&output.get_shape()) as usize;
+
+            results.insert(
+                output_name.to_string(),
+                bytemuck::cast_slice(&data)[..output_buffer_size].to_vec(),
+            );
             drop(data);
             buffer.unmap();
         }
diff --git a/src/optimisation.rs b/src/optimisation.rs
index 25be7a95..27df1280 100644
--- a/src/optimisation.rs
+++ b/src/optimisation.rs
@@ -1,8 +1,8 @@
 use crate::{
-    compiler::{compile, CompiledNode},
+    compiler::{compile, CompileError, CompiledNode},
     resource,
-    sequencer::sequence,
-    utils::{ceil, dimensions_infos, initializers, len},
+    sequencer::{sequence, SequenceError},
+    utils::{buffer_len, ceil, dimensions_infos, initializers},
     Result,
 };
 
@@ -10,6 +10,7 @@ use std::{borrow::Cow, collections::HashMap};
 
 use log::info;
 use tera::Tera;
+use thiserror::Error;
 use wgpu::BufferUsages;
 
 pub struct EncoderBuilder {
@@ -41,6 +42,11 @@ lazy_static! {
             include_str!("../templates/endomorphism/copy.wgsl"),
         )
         .unwrap();
+        tera.add_raw_template(
+            "endomorphism/softmax.wgsl",
+            include_str!("../templates/endomorphism/softmax.wgsl"),
+        )
+        .unwrap();
         tera.add_raw_template(
             "endomorphism/map.wgsl",
             include_str!("../templates/endomorphism/map.wgsl"),
@@ -114,10 +120,26 @@ lazy_static! {
 
 const MAX_BINDINGS_PER_GROUP: usize = 4;
 
+#[derive(Error, Debug)]
+pub enum OptimizationError {
+    #[error("compilation failed: {0}")]
+    CompilationFailed(#[from] CompileError),
+
+    #[error("output dims was not provided. You can use python's onnx-simplifier to generate implied dimensions.")]
+    OutputDimsMissing,
+
+    #[error("tensor metadata missing: '{0}'")]
+    TensorMetadataMissing(String),
+
+    #[error("sequencing failed: {0}")]
+    SequencingFailed(#[from] SequenceError),
+}
+
 pub fn load(
     graph: &crate::onnx::GraphProto,
     device: &wgpu::Device,
-) -> Result<(HashMap<String, wgpu::Buffer>, Vec<EncoderBuilder>)> {
+    opset_version: i64,
+) -> Result<(HashMap<String, wgpu::Buffer>, Vec<EncoderBuilder>), OptimizationError> {
     let initializers = initializers(graph);
     let dims_info = dimensions_infos(graph);
 
@@ -128,7 +150,7 @@ pub fn load(
             input_name.clone(),
             resource::buffer(
                 device,
-                len(input_dims) as _,
+                buffer_len(input_dims) as _,
                 input_name,
                 BufferUsages::STORAGE | BufferUsages::COPY_DST,
             ),
@@ -153,8 +175,9 @@ pub fn load(
 
         // Generate the shader source code for this node
         let (current_node, optimisation_length) =
-            sequence(&names, nodes, device, &initializers, &mut buffers);
-        let CompiledNode { shader, threads } = compile(&current_node, &dims_info, &TEMPLATES);
+            sequence(&names, nodes, device, &initializers, &mut buffers)?;
+        let CompiledNode { shader, threads } =
+            compile(&current_node, &dims_info, &TEMPLATES, opset_version)?;
         info!("shader: {}", shader);
 
         // Create buffers for all outputs of this node
@@ -168,7 +191,7 @@ pub fn load(
                         output.clone(),
                         resource::buffer(
                             device,
-                            len(output_dims) as _,
+                            buffer_len(output_dims) as _,
                             output.as_str(),
                             BufferUsages::STORAGE | BufferUsages::MAP_READ,
                         ),
@@ -178,14 +201,14 @@ pub fn load(
                         output.clone(),
                         resource::buffer(
                             device,
-                            len(output_dims) as _,
+                            buffer_len(output_dims) as _,
                             output.as_str(),
                             BufferUsages::STORAGE,
                         ),
                     );
                 }
             } else {
-                panic!("output dims was not provided. You can use python's onnx-simplifier to generate implied dimensions.")
+                return Err(OptimizationError::OutputDimsMissing);
             }
         }
 
@@ -200,9 +223,7 @@ pub fn load(
                 binding: binding_index,
                 resource: buffers
                     .get(tensor.as_str())
-                    .unwrap_or_else(|| {
-                        panic!("Tensor {} is not present in the inner infos", tensor)
-                    })
+                    .ok_or_else(|| OptimizationError::TensorMetadataMissing(tensor.to_string()))?
                     .as_entire_binding(),
             });
             binding_counter += 1;
@@ -215,9 +236,7 @@ pub fn load(
                 binding: binding_index,
                 resource: buffers
                     .get(tensor.as_str())
-                    .unwrap_or_else(|| {
-                        panic!("Tensor {} is not present in the inner infos", tensor)
-                    })
+                    .ok_or_else(|| OptimizationError::TensorMetadataMissing(tensor.to_string()))?
                     .as_entire_binding(),
             });
             binding_counter += 1;
diff --git a/src/sequencer.rs b/src/sequencer.rs
index a20c9f92..10ac7526 100644
--- a/src/sequencer.rs
+++ b/src/sequencer.rs
@@ -2,21 +2,31 @@ use std::collections::HashMap;
 
 use bytemuck::cast_slice;
 use protobuf::RepeatedField;
+use thiserror::Error;
 use wgpu::{Buffer, BufferUsages, Device};
 
 use crate::{
     onnx::NodeProto,
     resource::{self, padding},
-    utils::{attribute, get_attribute, node},
+    utils::{attribute, get_attribute, node, AttributeNotFoundError},
 };
 
+#[derive(Error, Debug)]
+pub enum SequenceError {
+    #[error("a required attribute was not found: {0}")]
+    AttributeNotFound(#[from] AttributeNotFoundError),
+
+    #[error("{0} is not implemented yet")]
+    NotImplemented(String),
+}
+
 pub fn sequence(
     names: &[&str],
     nodes: &[NodeProto],
     device: &Device,
     initializers: &HashMap<String, &[u8]>,
     inner_infos: &mut HashMap<String, Buffer>,
-) -> (NodeProto, usize) {
+) -> Result<(NodeProto, usize), SequenceError> {
     let mut optimisation_length = 1;
     let inputs = nodes[0].get_input();
 
@@ -36,12 +46,15 @@ pub fn sequence(
             for input in inputs {
                 if let Some(data) = initializers.get(input) {
                     let data = if input == &inputs[1]
-                        && get_attribute::<Vec<i64>>("kernel_shape", None, &nodes[0]) == [3, 3]
-                        && (get_attribute("pads", Some(vec![0, 0, 0, 0]), &nodes[0])
+                        && get_attribute::<Vec<i64>>("kernel_shape", None, &nodes[0])? == [3, 3]
+                        && (get_attribute("pads", Some(vec![0, 0, 0, 0]), &nodes[0])?
                             == [1, 1, 1, 1]
-                            || get_attribute("auto_pad", Some("SAME_UPPER".to_string()), &nodes[0])
-                                == "SAME_UPPER")
-                        && get_attribute("strides", Some(vec![1, 1]), &nodes[0]) == [1, 1]
+                            || get_attribute(
+                                "auto_pad",
+                                Some("SAME_UPPER".to_string()),
+                                &nodes[0],
+                            )? == "SAME_UPPER")
+                        && get_attribute("strides", Some(vec![1, 1]), &nodes[0])? == [1, 1]
                     {
                         padding(data, 12, 4)
                         // data.to_vec()
@@ -136,7 +149,9 @@ pub fn sequence(
                             if data == &[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] {
                                 attributes.push(attribute("coefficient", 0));
                             } else {
-                                unimplemented!()
+                                return Err(SequenceError::NotImplemented(String::from(
+                                    "Add with non-zero data",
+                                )));
                             }
                         }
                         _ => {
@@ -176,5 +191,5 @@ pub fn sequence(
         }
     };
 
-    (result, optimisation_length)
+    Ok((result, optimisation_length))
 }
diff --git a/src/utils.rs b/src/utils.rs
index c4caa05a..76638cf0 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,31 +1,50 @@
+use protobuf::RepeatedField;
+
 use crate::onnx;
+use crate::onnx::OperatorSetIdProto;
+use crate::onnx::ValueInfoProto;
 use std::collections::HashMap;
 use std::convert::From;
 use std::convert::Into;
 use std::str::from_utf8;
+use thiserror::Error;
+
+/* Minimum size of a buffer you can create with wgpu. Creating buffers smaller than this leads to panic "Validation
+* error: buffer binding size X is less than minimum 64" in Device::create_bind_group */
+const MINIMUM_BUFFER_SIZE: i64 = 64;
 
 pub fn len(dims: &[i64]) -> i64 {
     dims.iter().product::<i64>()
 }
 
+pub fn buffer_len(dims: &[i64]) -> i64 {
+    len(dims).max(MINIMUM_BUFFER_SIZE)
+}
+
+#[derive(Error, Debug)]
+#[error("did not find attribute '{attribute}' for node '{node_name}'")]
+pub struct AttributeNotFoundError {
+    attribute: String,
+    node_name: String,
+}
+
 pub fn get_attribute<T: std::convert::From<onnx::AttributeProto>>(
     attribute: &str,
     default: Option<T>,
     node: &onnx::NodeProto,
-) -> T {
+) -> Result<T, AttributeNotFoundError> {
     match (
         node.get_attribute()
             .iter()
             .find(|attr| attr.get_name() == attribute),
         default,
     ) {
-        (Some(attr), _) => attr.clone().into(),
-        (None, Some(default_attr)) => default_attr,
-        (None, None) => panic!(
-            "Did not find attribute: {} for node: {}",
-            attribute,
-            node.get_name()
-        ),
+        (Some(attr), _) => Ok(attr.clone().into()),
+        (None, Some(default_attr)) => Ok(default_attr),
+        (None, None) => Err(AttributeNotFoundError {
+            attribute: attribute.to_string(),
+            node_name: node.get_name().to_string(),
+        }),
     }
 }
 
@@ -42,42 +61,33 @@ pub fn rename_attribute(
     attr
 }
 
-pub fn dimensions_infos(graph_proto: &onnx::GraphProto) -> HashMap<String, Vec<i64>> {
-    let mut dims_info = HashMap::new();
-
-    for info in graph_proto.get_input() {
-        let dims = info
-            .get_field_type()
+impl ValueInfoProto {
+    pub fn get_shape(&self) -> Vec<i64> {
+        self.get_field_type()
             .get_tensor_type()
             .get_shape()
             .get_dim()
             .iter()
             .map(|x| x.get_dim_value())
-            .collect::<Vec<i64>>();
+            .collect::<Vec<i64>>()
+    }
+}
+
+pub fn dimensions_infos(graph_proto: &onnx::GraphProto) -> HashMap<String, Vec<i64>> {
+    let mut dims_info = HashMap::new();
+
+    for info in graph_proto.get_input() {
+        let dims = info.get_shape();
         dims_info.insert(info.get_name().to_string(), dims);
     }
 
     for info in graph_proto.get_output() {
-        let dims = info
-            .get_field_type()
-            .get_tensor_type()
-            .get_shape()
-            .get_dim()
-            .iter()
-            .map(|x| x.get_dim_value())
-            .collect::<Vec<i64>>();
+        let dims = info.get_shape();
         dims_info.insert(info.get_name().to_string(), dims);
     }
 
     for info in graph_proto.get_value_info() {
-        let dims = info
-            .get_field_type()
-            .get_tensor_type()
-            .get_shape()
-            .get_dim()
-            .iter()
-            .map(|x| x.get_dim_value())
-            .collect::<Vec<i64>>();
+        let dims = info.get_shape();
         dims_info.insert(info.get_name().to_string(), dims);
     }
 
@@ -189,6 +199,10 @@ pub fn graph(
 
 pub fn model(graph: onnx::GraphProto) -> onnx::ModelProto {
     let mut model = crate::onnx::ModelProto::new();
+    let mut onnx_opset_import = OperatorSetIdProto::new();
+    onnx_opset_import.set_domain("".to_string());
+    onnx_opset_import.set_version(13);
+    model.set_opset_import(RepeatedField::from_slice(&[onnx_opset_import]));
     model.set_graph(graph);
     model
 }
diff --git a/templates/endomorphism/softmax.wgsl b/templates/endomorphism/softmax.wgsl
new file mode 100644
index 00000000..7b327226
--- /dev/null
+++ b/templates/endomorphism/softmax.wgsl
@@ -0,0 +1,39 @@
+
+{%- include "structs.wgsl" -%}
+
+[[group(0), binding(0)]]
+var<storage, read> {{ inputs[0] }}: Array;
+
+[[group(0), binding(1)]]
+var<storage, write> {{ outputs[0] }}: Array;
+
+[[stage(compute), workgroup_size(1)]]
+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 = {{ inputs[0] }}.data[gidx + k];
+		max_element = max(max_element, element);
+	}
+
+	// Calculate sum(exp(input - max(input)))
+	var sum: f32 = 0.0;
+	for(var k: u32 = 0u; k < {{ i_lens[0] }}u; k = k + 1u) {
+		let element = {{ inputs[0] }}.data[gidx + k];
+		sum  = sum + exp(element - max_element);
+	}
+
+	// Calculate elements and write to output
+	for(var k: u32 = 0u; k < {{ i_lens[0] }}u; k = k + 1u) {
+		let element = {{ inputs[0] }}.data[gidx + k];
+		{{ outputs[0] }}.data[gidx + k] = exp(element - max_element) / sum;
+	}
+	{% endif %}
+}
\ No newline at end of file


Test patch:
diff --git a/tests/matrix.rs b/tests/matrix.rs
index 244187aa..bdc415b0 100644
--- a/tests/matrix.rs
+++ b/tests/matrix.rs
@@ -141,7 +141,7 @@ fn test_resize() {
         .expect("session did not create");
     let result = pollster::block_on(session.run(input_data)).unwrap();
 
-    let test_y = vec![1., 3., 0., 0.];
+    let test_y = vec![1., 3.];
     assert_eq!(result["Y"], test_y);
 
     let mut input_data = HashMap::new();
diff --git a/tests/pretrained_models.rs b/tests/pretrained_models.rs
index 77565453..e596bdc5 100644
--- a/tests/pretrained_models.rs
+++ b/tests/pretrained_models.rs
@@ -17,7 +17,7 @@ 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, 0.0, 0.0]);
+    assert_eq!(result["y"], &[0.0, 1.0]);
 }
 
 #[test]