guess batch dimension when inference fails

7 files changed, 122 insertions, 7 deletions

diff --git a/eval/src/tests/tensor/onnx_wrapper/guess_batch.onnx b/eval/src/tests/tensor/onnx_wrapper/guess_batch.onnx
new file mode 100644
index 00000000000..9ae224406eb
--- /dev/null
+++ b/eval/src/tests/tensor/onnx_wrapper/guess_batch.onnx
@@ -0,0 +1,16 @@
+guess_batch.py:n
+
+in1
+in2out"Addguess_batchZ
+in1
+

+
+batch1Z
+in2
+

+
+batch2b
+out
+

+
+batch3B
\ No newline at end of file
diff --git a/eval/src/tests/tensor/onnx_wrapper/guess_batch.py b/eval/src/tests/tensor/onnx_wrapper/guess_batch.py
new file mode 100755
index 00000000000..c43448c58a7
--- /dev/null
+++ b/eval/src/tests/tensor/onnx_wrapper/guess_batch.py
@@ -0,0 +1,26 @@
+# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+import onnx
+from onnx import helper, TensorProto
+
+IN1 = helper.make_tensor_value_info('in1', TensorProto.FLOAT, ['batch1'])
+IN2 = helper.make_tensor_value_info('in2', TensorProto.FLOAT, ['batch2'])
+OUT = helper.make_tensor_value_info('out', TensorProto.FLOAT, ['batch3'])
+
+nodes = [
+    helper.make_node(
+        'Add',
+        ['in1', 'in2'],
+        ['out'],
+    ),
+]
+graph_def = helper.make_graph(
+    nodes,
+    'guess_batch',
+    [
+        IN1,
+        IN2,
+    ],
+    [OUT],
+)
+model_def = helper.make_model(graph_def, producer_name='guess_batch.py')
+onnx.save(model_def, 'guess_batch.onnx')
diff --git a/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp b/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp
index 23c41167266..c733f922194 100644
--- a/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp
+++ b/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp
@@ -22,6 +22,7 @@ std::string source_dir = get_source_dir();
 std::string simple_model = source_dir + "/simple.onnx";
 std::string dynamic_model = source_dir + "/dynamic.onnx";
 std::string int_types_model = source_dir + "/int_types.onnx";
+std::string guess_batch_model = source_dir + "/guess_batch.onnx";
 
 void dump_info(const char *ctx, const std::vector<TensorInfo> &info) {
     fprintf(stderr, "%s:\n", ctx);
@@ -266,4 +267,45 @@ TEST(OnnxTest, int_types_onnx_model_can_be_evaluated)
     //-------------------------------------------------------------------------
 }
 
+TEST(OnnxTest, we_guess_batch_dimension_size_when_inference_fails) {    
+    Onnx model(guess_batch_model, Onnx::Optimize::ENABLE);
+    Onnx::WirePlanner planner_3;
+    Onnx::WirePlanner planner_4;
+
+    ValueType in_3_type = ValueType::from_spec("tensor<float>(a[3])");
+    std::vector<float> in_3_values({1.0, 2.0, 3.0});
+    DenseTensorView in_3(in_3_type, TypedCells(in_3_values));
+    EXPECT_TRUE(planner_3.bind_input_type(in_3_type, model.inputs()[0]));
+    EXPECT_TRUE(planner_3.bind_input_type(in_3_type, model.inputs()[1]));
+
+    ValueType in_4_type = ValueType::from_spec("tensor<float>(a[4])");
+    std::vector<float> in_4_values({1.0, 2.0, 3.0, 4.0});
+    DenseTensorView in_4(in_4_type, TypedCells(in_4_values));
+    EXPECT_TRUE(planner_4.bind_input_type(in_4_type, model.inputs()[0]));
+    EXPECT_TRUE(planner_4.bind_input_type(in_4_type, model.inputs()[1]));
+
+    EXPECT_EQ(planner_3.make_output_type(model.outputs()[0]).to_spec(), "tensor<float>(d0[3])");
+    EXPECT_EQ(planner_4.make_output_type(model.outputs()[0]).to_spec(), "tensor<float>(d0[4])");
+
+    Onnx::WireInfo wire_info_3 = planner_3.get_wire_info(model);
+    Onnx::WireInfo wire_info_4 = planner_4.get_wire_info(model);
+    Onnx::EvalContext ctx_3(model, wire_info_3);
+    Onnx::EvalContext ctx_4(model, wire_info_4);
+    //-------------------------------------------------------------------------
+    ctx_3.bind_param(0, in_3);
+    ctx_3.bind_param(1, in_3);
+    ctx_3.eval();
+    ctx_4.bind_param(0, in_4);
+    ctx_4.bind_param(1, in_4);
+    ctx_4.eval();
+    //-------------------------------------------------------------------------
+    auto out_3 = TensorSpec::from_value(ctx_3.get_result(0));
+    auto out_4 = TensorSpec::from_value(ctx_4.get_result(0));
+    auto expect_3 = TensorSpec::from_expr("tensor<float>(d0[3]):[2,4,6]");
+    auto expect_4 = TensorSpec::from_expr("tensor<float>(d0[4]):[2,4,6,8]");
+    EXPECT_EQ(out_3, expect_3);
+    EXPECT_EQ(out_4, expect_4);
+    //-------------------------------------------------------------------------
+}
+
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/eval/src/vespa/eval/eval/tensor_spec.cpp b/eval/src/vespa/eval/eval/tensor_spec.cpp
index b4b2e3d3afc..712ee282c71 100644
--- a/eval/src/vespa/eval/eval/tensor_spec.cpp
+++ b/eval/src/vespa/eval/eval/tensor_spec.cpp
@@ -4,6 +4,9 @@
 #include "value.h"
 #include "tensor.h"
 #include "tensor_engine.h"
+#include "simple_tensor_engine.h"
+#include "function.h"
+#include "interpreted_function.h"
 #include <vespa/vespalib/util/stringfmt.h>
 #include <vespa/vespalib/data/slime/slime.h>
 #include <ostream>
@@ -109,6 +112,19 @@ TensorSpec::from_value(const eval::Value &value)
     return TensorSpec("error");
 }
 
+TensorSpec
+TensorSpec::from_expr(const vespalib::string &expr)
+{
+    NoParams no_params;
+    auto fun = Function::parse(expr);
+    if (!fun->has_error() && (fun->num_params() == 0)) {
+        InterpretedFunction ifun(SimpleTensorEngine::ref(), *fun, NodeTypes());
+        InterpretedFunction::Context ctx(ifun);
+        return from_value(ifun.eval(ctx, no_params));
+    }
+    return TensorSpec("error");
+}
+
 bool
 operator==(const TensorSpec &lhs, const TensorSpec &rhs)
 {
diff --git a/eval/src/vespa/eval/eval/tensor_spec.h b/eval/src/vespa/eval/eval/tensor_spec.h
index 06f7c39a2d3..974ad4a1f4c 100644
--- a/eval/src/vespa/eval/eval/tensor_spec.h
+++ b/eval/src/vespa/eval/eval/tensor_spec.h
@@ -81,6 +81,7 @@ public:
     void to_slime(slime::Cursor &tensor) const;
     static TensorSpec from_slime(const slime::Inspector &tensor);
     static TensorSpec from_value(const eval::Value &value);
+    static TensorSpec from_expr(const vespalib::string &expr);
 };
 
 bool operator==(const TensorSpec &lhs, const TensorSpec &rhs);
diff --git a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp
index 8f1b01a58ab..0e545bf95c3 100644
--- a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp
+++ b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp
@@ -256,12 +256,15 @@ Onnx::WirePlanner::bind_input_type(const eval::ValueType &vespa_in, const Tensor
             if (dimensions[i].value != type.dimensions()[i].size) {
                 return false;
             }
-        } else if (dimensions[i].is_symbolic()) {
-            auto &bound_size = _symbolic_sizes[dimensions[i].name];
-            if (bound_size == 0) {
-                bound_size = type.dimensions()[i].size;
-            } else if (bound_size != type.dimensions()[i].size) {
-                return false;
+        } else {
+            _unknown_sizes.insert(type.dimensions()[i].size);
+            if (dimensions[i].is_symbolic()) {
+                auto &bound_size = _symbolic_sizes[dimensions[i].name];
+                if (bound_size == 0) {
+                    bound_size = type.dimensions()[i].size;
+                } else if (bound_size != type.dimensions()[i].size) {
+                    return false;
+                }
             }
         }
     }
@@ -283,6 +286,15 @@ Onnx::WirePlanner::make_output_type(const TensorInfo &onnx_out) const
                 dim_size = pos->second;
             }
         }
+        // if the output dimension is still unknown, but all unknown
+        // input dimensions have the same size, we use that size as a
+        // guess for the size of the unknown output dimension as
+        // well. (typical scenario would be batch dimension not tagged
+        // as having the same symbolic size across input and output
+        // values).
+        if ((dim_size == 0) && (_unknown_sizes.size() == 1)) {
+            dim_size = *_unknown_sizes.begin();
+        }
         if ((dim_size == 0) || (dim_list.size() > 9)) {
             return ValueType::error_type();
         }
diff --git a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h
index 85a59824229..a059f2e2494 100644
--- a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h
+++ b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h
@@ -12,6 +12,7 @@
 #include <vespa/eval/eval/value_type.h>
 #include <vector>
 #include <map>
+#include <set>
 
 namespace vespalib::eval { struct Value; }
 
@@ -83,8 +84,9 @@ public:
     private:
         std::map<vespalib::string,eval::ValueType> _input_types;
         std::map<vespalib::string,size_t> _symbolic_sizes;
+        std::set<size_t> _unknown_sizes;
     public:
-        WirePlanner() : _input_types(), _symbolic_sizes() {}
+        WirePlanner() : _input_types(), _symbolic_sizes(), _unknown_sizes() {}
         ~WirePlanner();
         bool bind_input_type(const eval::ValueType &vespa_in, const TensorInfo &onnx_in);
         eval::ValueType make_output_type(const TensorInfo &onnx_out) const;