fix for ? in onnx shapes specs

sonos · Apr 22, 2024 · 7dfcdc9 · 7dfcdc9
1 parent dae9667
commit 7dfcdc9
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/onnx/src/model.rs b/onnx/src/model.rs
@@ -6,9 +6,9 @@ use std::collections::HashMap;
 use tract_hir::internal::*;
 use tract_hir::prelude::tract_itertools::Itertools;
 
+use crate::data_resolver::{self, ModelDataResolver};
 use crate::pb::type_proto::Value;
 use crate::pb::{self, TensorProto, TypeProto};
-use crate::data_resolver::{self, ModelDataResolver};
 use crate::tensor::{load_tensor, translate_inference_fact};
 use prost::Message;
 
@@ -88,7 +88,8 @@ impl<'a> ParsingContext<'a> {
                 let fact = input.r#type.as_ref().unwrap().value.as_ref().unwrap();
                 #[allow(irrefutable_let_patterns)]
                 let fact: InferenceFact = if let pb::type_proto::Value::TensorType(fact) = fact {
-                    translate_inference_fact(&ctx, fact, true)?
+                    translate_inference_fact(&ctx, fact, true)
+                        .with_context(|| format!("translating to fact: {:?}", fact))?
                 } else {
                     bail!("Can not parse tensor type");
                 };

diff --git a/onnx/src/tensor.rs b/onnx/src/tensor.rs
@@ -1,7 +1,7 @@
+use crate::data_resolver::ModelDataResolver;
 use crate::model::ParsingContext;
 use crate::pb::tensor_proto::DataType;
 use crate::pb::*;
-use crate::data_resolver::ModelDataResolver;
 use prost::Message;
 use std::convert::{TryFrom, TryInto};
 use std::path::PathBuf;
@@ -46,10 +46,11 @@ pub fn translate_inference_fact(
                         Ok(DimFact::from(v.to_dim()))
                     }
                     Some(tensor_shape_proto::dimension::Value::DimParam(v)) => {
-                        if v.starts_with("unk__") && !include_unknown_symbols {
+                        if v == "?" || (v.starts_with("unk__") && !include_unknown_symbols) {
                             Ok(DimFact::default())
                         } else {
-                            let dim = parse_tdim(&ctx.symbol_table, v)?;
+                            let dim = parse_tdim(&ctx.symbol_table, v)
+                                .with_context(|| format!("Parsing as TDim: `{v}'"))?;
                             Ok(DimFact::from(dim))
                         }
                     }