Merge branch 'master' into bratseth/typecheck-all

4 files changed, 33 insertions, 12 deletions

diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/OperationMapper.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/OperationMapper.java
index b8f8e288257..55782c36d18 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/OperationMapper.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/OperationMapper.java
@@ -219,7 +219,7 @@ class OperationMapper {
     private static Optional<TypedTensorFunction> placeholderWithDefault(TensorFlowImporter.Parameters params) {
         String name = toVespaName(params.node().getInput(0));
         Tensor defaultValue = getConstantTensor(params, params.node().getInput(0));
-        params.result().constant(name, defaultValue);
+        params.result().largeConstant(name, defaultValue);
         params.result().macro(name, new RankingExpression(name, new ReferenceNode("constant(\"" + name + "\")")));
         // The default value will be provided by the macro. Users can override macro to change value.
         TypedTensorFunction output = new TypedTensorFunction(defaultValue.type(), new VariableTensor(name));
@@ -544,7 +544,11 @@ class OperationMapper {
 
     private static Optional<TypedTensorFunction> createConstant(TensorFlowImporter.Parameters params, Tensor constant) {
         String name = toVespaName(params.node().getName());
-        params.result().constant(name, constant);
+        if (constant.type().rank() == 0 || constant.size() <= 1) {
+            params.result().smallConstant(name, constant);
+        } else {
+            params.result().largeConstant(name, constant);
+        }
         TypedTensorFunction output = new TypedTensorFunction(constant.type(),
                 new TensorFunctionNode.TensorFunctionExpressionNode(
                         new ReferenceNode("constant(\"" + name + "\")")));
@@ -553,8 +557,11 @@ class OperationMapper {
 
     private static Tensor getConstantTensor(TensorFlowImporter.Parameters params, String name) {
         String vespaName = toVespaName(name);
-        if (params.result().constants().containsKey(vespaName)) {
-            return params.result().constants().get(vespaName);
+        if (params.result().smallConstants().containsKey(vespaName)) {
+            return params.result().smallConstants().get(vespaName);
+        }
+        if (params.result().largeConstants().containsKey(vespaName)) {
+            return params.result().largeConstants().get(vespaName);
         }
         Session.Runner fetched = params.model().session().runner().fetch(name);
         List<org.tensorflow.Tensor<?>> importedTensors = fetched.run();
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowModel.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowModel.java
index 530f4793b62..351aa417f9c 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowModel.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowModel.java
@@ -24,13 +24,15 @@ public class TensorFlowModel {
 
     private final Map<String, Signature> signatures = new HashMap<>();
     private final Map<String, TensorType> arguments = new HashMap<>();
-    private final Map<String, Tensor> constants = 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> macros = new HashMap<>();
     private final Map<String, TensorType> requiredMacros = new HashMap<>();
 
     void argument(String name, TensorType argumentType) { arguments.put(name, argumentType); }
-    void constant(String name, Tensor constant) { constants.put(name, constant); }
+    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 macro(String name, RankingExpression expression) { macros.put(name, expression); }
     void requiredMacro(String name, TensorType type) { requiredMacros.put(name, type); }
@@ -43,8 +45,19 @@ public class TensorFlowModel {
     /** 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 constants of this */
-    public Map<String, Tensor> constants() { return Collections.unmodifiableMap(constants); }
+    /**
+     * Returns an immutable map of the small constants of this.
+     * These should have sizes up to a few kb at most, and correspond to constant
+     * values given in the TensorFlow source.
+     */
+    public Map<String, Tensor> smallConstants() { return Collections.unmodifiableMap(smallConstants); }
+
+    /**
+     * Returns an immutable map of the large constants of this.
+     * These can have sizes in gigabytes and must be distributed to nodes separately from configuration,
+     * and correspond to Variable files stored separately in TensorFlow.
+     */
+    public Map<String, Tensor> largeConstants() { return Collections.unmodifiableMap(largeConstants); }
 
     /**
      * Returns an immutable map of the expressions of this - corresponding to TensorFlow nodes
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/MnistSoftmaxImportTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/MnistSoftmaxImportTestCase.java
index 01dd15d5fa0..ad5abd4c03d 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/MnistSoftmaxImportTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/MnistSoftmaxImportTestCase.java
@@ -20,15 +20,15 @@ public class MnistSoftmaxImportTestCase {
         TestableTensorFlowModel model = new TestableTensorFlowModel("src/test/files/integration/tensorflow/mnist_softmax/saved");
 
         // Check constants
-        assertEquals(2, model.get().constants().size());
+        assertEquals(2, model.get().largeConstants().size());
 
-        Tensor constant0 = model.get().constants().get("Variable");
+        Tensor constant0 = model.get().largeConstants().get("Variable");
         assertNotNull(constant0);
         assertEquals(new TensorType.Builder().indexed("d0", 784).indexed("d1", 10).build(),
                      constant0.type());
         assertEquals(7840, constant0.size());
 
-        Tensor constant1 = model.get().constants().get("Variable_1");
+        Tensor constant1 = model.get().largeConstants().get("Variable_1");
         assertNotNull(constant1);
         assertEquals(new TensorType.Builder().indexed("d0", 10).build(),
                      constant1.type());
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TestableTensorFlowModel.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TestableTensorFlowModel.java
index 2c621fd2e92..ae7714b271a 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TestableTensorFlowModel.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TestableTensorFlowModel.java
@@ -57,7 +57,8 @@ public class TestableTensorFlowModel {
 
     private Context contextFrom(TensorFlowModel result) {
         MapContext context = new MapContext();
-        result.constants().forEach((name, tensor) -> context.put("constant(\"" + name + "\")", new TensorValue(tensor)));
+        result.largeConstants().forEach((name, tensor) -> context.put("constant(\"" + name + "\")", new TensorValue(tensor)));
+        result.smallConstants().forEach((name, tensor) -> context.put("constant(\"" + name + "\")", new TensorValue(tensor)));
         return context;
     }