Propagate input type information

8 files changed, 94 insertions, 57 deletions

diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/ImportedModel.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/ImportedModel.java
index 282a4c5e0a9..88b5645e2e5 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/ImportedModel.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/ImportedModel.java
@@ -1,15 +1,19 @@
 package com.yahoo.searchlib.rankingexpression.integration.ml;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.yahoo.collections.Pair;
 import com.yahoo.searchlib.rankingexpression.RankingExpression;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorType;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.regex.Pattern;
 
 /**
@@ -26,12 +30,11 @@ public class ImportedModel {
     private final String source;
 
     private final Map<String, Signature> signatures = new HashMap<>();
-    private final Map<String, TensorType> arguments = new HashMap<>();
+    private final Map<String, TensorType> inputs = new HashMap<>();
     private final Map<String, Tensor> smallConstants = new HashMap<>();
     private final Map<String, Tensor> largeConstants = new HashMap<>();
     private final Map<String, RankingExpression> expressions = new HashMap<>();
     private final Map<String, RankingExpression> functions = new HashMap<>();
-    private final Map<String, TensorType> requiredFunctions = new HashMap<>();
 
     /**
      * Creates a new imported model.
@@ -49,11 +52,11 @@ public class ImportedModel {
     /** Returns the name of this model, which can only contain the characters in [A-Za-z0-9_] */
     public String name() { return name; }
 
-    /** Returns the source path (directiry or file) of this model */
+    /** Returns the source path (directory or file) of this model */
     public String source() { return source; }
 
-    /** Returns an immutable map of the arguments ("Placeholders") of this */
-    public Map<String, TensorType> arguments() { return Collections.unmodifiableMap(arguments); }
+    /** Returns an immutable map of the inputs of this */
+    public Map<String, TensorType> inputs() { return Collections.unmodifiableMap(inputs); }
 
     /**
      * Returns an immutable map of the small constants of this.
@@ -71,7 +74,7 @@ public class ImportedModel {
 
     /**
      * Returns an immutable map of the expressions of this - corresponding to graph nodes
-     * which are not Inputs/Placeholders or Variables (which instead become respectively arguments and constants).
+     * which are not Inputs/Placeholders or Variables (which instead become respectively inputs and constants).
      * Note that only nodes recursively referenced by a placeholder/input are added.
      */
     public Map<String, RankingExpression> expressions() { return Collections.unmodifiableMap(expressions); }
@@ -82,9 +85,6 @@ public class ImportedModel {
      */
     public Map<String, RankingExpression> functions() { return Collections.unmodifiableMap(functions); }
 
-    /** Returns an immutable map of the functions that must be provided by the environment running this model */
-    public Map<String, TensorType> requiredFunctions() { return Collections.unmodifiableMap(requiredFunctions); }
-
     /** Returns an immutable map of the signatures of this */
     public Map<String, Signature> signatures() { return Collections.unmodifiableMap(signatures); }
 
@@ -96,12 +96,11 @@ public class ImportedModel {
     /** Convenience method for returning a default signature */
     Signature defaultSignature() { return signature(defaultSignatureName); }
 
-    void argument(String name, TensorType argumentType) { arguments.put(name, argumentType); }
+    void input(String name, TensorType argumentType) { inputs.put(name, argumentType); }
     void smallConstant(String name, Tensor constant) { smallConstants.put(name, constant); }
     void largeConstant(String name, Tensor constant) { largeConstants.put(name, constant); }
     void expression(String name, RankingExpression expression) { expressions.put(name, expression); }
     void function(String name, RankingExpression expression) { functions.put(name, expression); }
-    void requiredFunction(String name, TensorType type) { requiredFunctions.put(name, type); }
 
     /**
      * Returns all the output expressions of this indexed by name. The names consist of one or two parts
@@ -109,24 +108,27 @@ public class ImportedModel {
      * if signatures are used, or the expression name if signatures are not used and there are multiple
      * expressions, and the second is the output name if signature names are used.
      */
-    public List<Pair<String, RankingExpression>> outputExpressions() {
-        List<Pair<String, RankingExpression>> expressions = new ArrayList<>();
+    public List<Pair<String, ExpressionWithInputs>> outputExpressions() {
+        List<Pair<String, ExpressionWithInputs>> expressions = new ArrayList<>();
         for (Map.Entry<String, Signature> signatureEntry : signatures().entrySet()) {
             for (Map.Entry<String, String> outputEntry : signatureEntry.getValue().outputs().entrySet())
                 expressions.add(new Pair<>(signatureEntry.getKey() + "." + outputEntry.getKey(),
-                                           expressions().get(outputEntry.getValue())));
+                                           signatureEntry.getValue().outputExpression(outputEntry.getKey())));
             if (signatureEntry.getValue().outputs().isEmpty()) // fallback: Signature without outputs
                 expressions.add(new Pair<>(signatureEntry.getKey(),
-                                           expressions().get(signatureEntry.getKey())));
+                                           new ExpressionWithInputs(expressions().get(signatureEntry.getKey()),
+                                                                    signatureEntry.getValue().inputMap())));
         }
         if (signatures().isEmpty()) { // fallback for models without signatures
             if (expressions().size() == 1) {
                 Map.Entry<String, RankingExpression> singleEntry = this.expressions.entrySet().iterator().next();
-                expressions.add(new Pair<>(singleEntry.getKey(), singleEntry.getValue()));
+                expressions.add(new Pair<>(singleEntry.getKey(),
+                                           new ExpressionWithInputs(singleEntry.getValue(), inputs)));
             }
             else {
                 for (Map.Entry<String, RankingExpression> expressionEntry : expressions().entrySet()) {
-                    expressions.add(new Pair<>(expressionEntry.getKey(), expressionEntry.getValue()));
+                    expressions.add(new Pair<>(expressionEntry.getKey(),
+                                               new ExpressionWithInputs(expressionEntry.getValue(), inputs)));
                 }
             }
         }
@@ -134,7 +136,7 @@ public class ImportedModel {
     }
 
     /**
-     * A signature is a set of named inputs and outputs, where the inputs maps to argument
+     * A signature is a set of named inputs and outputs, where the inputs maps to input
      * ("placeholder") names+types, and outputs maps to expressions nodes.
      * Note that TensorFlow supports multiple signatures in their format, but ONNX has no explicit
      * concept of signatures. For now, we handle ONNX models as having a single signature.
@@ -158,12 +160,20 @@ public class ImportedModel {
 
         /**
          * Returns an immutable map of the inputs (evaluation context) of this. This is a map from input name
-         * to argument (Placeholder) name in the owner of this
+         * in this signature to input name in the owning model
          */
         public Map<String, String> inputs() { return Collections.unmodifiableMap(inputs); }
 
-        /** Returns the type of the argument this input references */
-        public TensorType inputArgument(String inputName) { return owner().arguments().get(inputs.get(inputName)); }
+        /** Returns the name and type of all inputs in this signature as an immutable map */
+        public Map<String, TensorType> inputMap() {
+            ImmutableMap.Builder<String, TensorType> inputs = new ImmutableMap.Builder<>();
+            for (Map.Entry<String, String> inputEntry : inputs().entrySet())
+                inputs.put(inputEntry.getKey(), owner().inputs().get(inputEntry.getValue()));
+            return inputs.build();
+        }
+
+        /** Returns the type of the input this input references */
+        public TensorType inputArgument(String inputName) { return owner().inputs().get(inputs.get(inputName)); }
 
         /** Returns an immutable list of the expression names of this */
         public Map<String, String> outputs() { return Collections.unmodifiableMap(outputs); }
@@ -180,7 +190,9 @@ public class ImportedModel {
         public List<String> importWarnings() { return Collections.unmodifiableList(importWarnings); }
 
         /** Returns the expression this output references */
-        public RankingExpression outputExpression(String outputName) { return owner().expressions().get(outputs.get(outputName)); }
+        public ExpressionWithInputs outputExpression(String outputName) {
+            return new ExpressionWithInputs(owner().expressions().get(outputs.get(outputName)), inputMap());
+        }
 
         @Override
         public String toString() { return "signature '" + name + "'"; }
@@ -192,4 +204,28 @@ public class ImportedModel {
 
     }
 
+    /**
+     * An expression, with the inputs (bindings) which must be supplied to evaluate it.
+     * All non-scalar (non-empty tensor type) inputs are always present here. Inputs not
+     * given explicitly here (but present in the expression) are always scalar.
+     */
+    public static class ExpressionWithInputs {
+
+        private final RankingExpression expression;
+        private final ImmutableMap<String, TensorType> inputs;
+
+        public ExpressionWithInputs(RankingExpression expression, Map<String, TensorType> inputs) {
+            this.expression = Objects.requireNonNull(expression, "expression cannot be null");
+            this.inputs = ImmutableMap.copyOf(inputs);
+        }
+
+        public RankingExpression expression() { return expression; }
+        public ImmutableMap<String, TensorType> inputs() { return inputs; }
+
+        public ExpressionWithInputs with(RankingExpression newExpression) {
+            return new ExpressionWithInputs(newExpression, inputs);
+        }
+
+    }
+
 }
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/ModelImporter.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/ModelImporter.java
index d25502fd149..b7138ad87e3 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/ModelImporter.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/ModelImporter.java
@@ -187,8 +187,7 @@ public abstract class ModelImporter {
         if (operation.isInput()) {
             // All inputs must have dimensions with standard naming convention: d0, d1, ...
             OrderedTensorType standardNamingConvention = OrderedTensorType.standardType(operation.type().get());
-            model.argument(operation.vespaName(), standardNamingConvention.type());
-            model.requiredFunction(operation.vespaName(), standardNamingConvention.type());
+            model.input(operation.vespaName(), standardNamingConvention.type());
         }
     }
 
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/OnnxImporter.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/OnnxImporter.java
index 917b0d6a389..e6bb5f40b3f 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/OnnxImporter.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/ml/OnnxImporter.java
@@ -2,7 +2,6 @@
 
 package com.yahoo.searchlib.rankingexpression.integration.ml;
 
-import com.yahoo.io.IOUtils;
 import com.yahoo.searchlib.rankingexpression.integration.ml.importer.IntermediateGraph;
 import com.yahoo.searchlib.rankingexpression.integration.ml.importer.onnx.GraphImporter;
 import onnx.Onnx;
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/BatchNormImportTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/BatchNormImportTestCase.java
index bf9684082f4..3a1c9ec9551 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/BatchNormImportTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/BatchNormImportTestCase.java
@@ -20,10 +20,11 @@ public class BatchNormImportTestCase {
         assertEquals("Has skipped outputs",
                      0, model.get().signature("serving_default").skippedOutputs().size());
 
-        RankingExpression output = signature.outputExpression("y");
+        ImportedModel.ExpressionWithInputs output = signature.outputExpression("y");
         assertNotNull(output);
-        assertEquals("dnn/batch_normalization_3/batchnorm/add_1", output.getName());
-        model.assertEqualResult("X", output.getName());
+        assertEquals("dnn/batch_normalization_3/batchnorm/add_1", output.expression().getName());
+        model.assertEqualResult("X", output.expression().getName());
+        assertEquals("{x=tensor(d0[],d1[784])}", output.inputs().toString());
     }
 
 }
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/DropoutImportTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/DropoutImportTestCase.java
index a8f7542f3a4..4c35d843f5d 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/DropoutImportTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/DropoutImportTestCase.java
@@ -19,22 +19,23 @@ public class DropoutImportTestCase {
         TestableTensorFlowModel model = new TestableTensorFlowModel("test", "src/test/files/integration/tensorflow/dropout/saved");
 
         // Check required functions
-        assertEquals(1, model.get().requiredFunctions().size());
-        assertTrue(model.get().requiredFunctions().containsKey("X"));
+        assertEquals(1, model.get().inputs().size());
+        assertTrue(model.get().inputs().containsKey("X"));
         assertEquals(new TensorType.Builder().indexed("d0").indexed("d1", 784).build(),
-                     model.get().requiredFunctions().get("X"));
+                     model.get().inputs().get("X"));
 
         ImportedModel.Signature signature = model.get().signature("serving_default");
 
         assertEquals("Has skipped outputs",
                      0, model.get().signature("serving_default").skippedOutputs().size());
 
-        RankingExpression output = signature.outputExpression("y");
+        ImportedModel.ExpressionWithInputs output = signature.outputExpression("y");
         assertNotNull(output);
-        assertEquals("outputs/Maximum", output.getName());
+        assertEquals("outputs/Maximum", output.expression().getName());
         assertEquals("join(join(imported_ml_function_test_outputs_BiasAdd, reduce(constant(test_outputs_Const), sum, d1), f(a,b)(a * b)), imported_ml_function_test_outputs_BiasAdd, f(a,b)(max(a,b)))",
-                output.getRoot().toString());
-        model.assertEqualResult("X", output.getName());
+                     output.expression().getRoot().toString());
+        model.assertEqualResult("X", output.expression().getName());
+        assertEquals("{x=tensor(d0[],d1[784])}", output.inputs().toString());
     }
 
 }
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/MnistImportTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/MnistImportTestCase.java
index add66eece1a..b3e281ad25d 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/MnistImportTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/MnistImportTestCase.java
@@ -20,11 +20,10 @@ public class MnistImportTestCase {
         assertEquals("Has skipped outputs",
                      0, model.get().signature("serving_default").skippedOutputs().size());
 
-        RankingExpression output = signature.outputExpression("y");
+        ImportedModel.ExpressionWithInputs output = signature.outputExpression("y");
         assertNotNull(output);
-        assertEquals("dnn/outputs/add", output.getName());
-        model.assertEqualResultSum("input", output.getName(), 0.00001);
+        assertEquals("dnn/outputs/add", output.expression().getName());
+        model.assertEqualResultSum("input", output.expression().getName(), 0.00001);
     }
 
-
 }
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/OnnxMnistSoftmaxImportTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/OnnxMnistSoftmaxImportTestCase.java
index e20ac16a691..b5655cfbfa5 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/OnnxMnistSoftmaxImportTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/OnnxMnistSoftmaxImportTestCase.java
@@ -27,27 +27,28 @@ public class OnnxMnistSoftmaxImportTestCase {
         Tensor constant0 = model.largeConstants().get("test_Variable");
         assertNotNull(constant0);
         assertEquals(new TensorType.Builder().indexed("d2", 784).indexed("d1", 10).build(),
-                constant0.type());
+                     constant0.type());
         assertEquals(7840, constant0.size());
 
         Tensor constant1 = model.largeConstants().get("test_Variable_1");
         assertNotNull(constant1);
-        assertEquals(new TensorType.Builder().indexed("d1", 10).build(),
-                constant1.type());
+        assertEquals(new TensorType.Builder().indexed("d1", 10).build(), constant1.type());
         assertEquals(10, constant1.size());
 
-        // Check required functions (inputs)
-        assertEquals(1, model.requiredFunctions().size());
-        assertTrue(model.requiredFunctions().containsKey("Placeholder"));
-        assertEquals(new TensorType.Builder().indexed("d0").indexed("d1", 784).build(),
-                     model.requiredFunctions().get("Placeholder"));
+        // Check inputs
+        assertEquals(1, model.inputs().size());
+        assertTrue(model.inputs().containsKey("Placeholder"));
+        assertEquals(TensorType.fromSpec("tensor(d0[],d1[784])"), model.inputs().get("Placeholder"));
 
-        // Check outputs
-        RankingExpression output = model.defaultSignature().outputExpression("add");
+        // Check signature
+        ImportedModel.ExpressionWithInputs output = model.defaultSignature().outputExpression("add");
         assertNotNull(output);
-        assertEquals("add", output.getName());
+        assertEquals("add", output.expression().getName());
         assertEquals("join(reduce(join(rename(Placeholder, (d0, d1), (d0, d2)), constant(test_Variable), f(a,b)(a * b)), sum, d2), constant(test_Variable_1), f(a,b)(a + b))",
-                     output.getRoot().toString());
+                     output.expression().getRoot().toString());
+        assertEquals(TensorType.fromSpec("tensor(d0[],d1[784])"),
+                     model.inputs().get(model.defaultSignature().inputs().get("Placeholder")));
+        assertEquals("{Placeholder=tensor(d0[],d1[784])}", output.inputs().toString());
     }
 
     @Test
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/TensorFlowMnistSoftmaxImportTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/TensorFlowMnistSoftmaxImportTestCase.java
index ef28eb4678f..4a0362c0229 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/TensorFlowMnistSoftmaxImportTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/TensorFlowMnistSoftmaxImportTestCase.java
@@ -38,10 +38,10 @@ public class TensorFlowMnistSoftmaxImportTestCase {
         assertEquals(0, model.get().functions().size());
 
         // Check required functions
-        assertEquals(1, model.get().requiredFunctions().size());
-        assertTrue(model.get().requiredFunctions().containsKey("Placeholder"));
+        assertEquals(1, model.get().inputs().size());
+        assertTrue(model.get().inputs().containsKey("Placeholder"));
         assertEquals(new TensorType.Builder().indexed("d0").indexed("d1", 784).build(),
-                     model.get().requiredFunctions().get("Placeholder"));
+                     model.get().inputs().get("Placeholder"));
 
         // Check signatures
         assertEquals(1, model.get().signatures().size());
@@ -56,11 +56,12 @@ public class TensorFlowMnistSoftmaxImportTestCase {
 
         // ... signature outputs
         assertEquals(1, signature.outputs().size());
-        RankingExpression output = signature.outputExpression("y");
+        ImportedModel.ExpressionWithInputs output = signature.outputExpression("y");
         assertNotNull(output);
-        assertEquals("add", output.getName());
+        assertEquals("add", output.expression().getName());
         assertEquals("join(reduce(join(rename(Placeholder, (d0, d1), (d0, d2)), constant(test_Variable_read), f(a,b)(a * b)), sum, d2), constant(test_Variable_1_read), f(a,b)(a + b))",
-                     output.getRoot().toString());
+                     output.expression().getRoot().toString());
+        assertEquals("{x=tensor(d0[],d1[784])}", output.inputs().toString());
 
         // Test execution
         model.assertEqualResult("Placeholder", "MatMul");