Generate macros for TensorFlow nodes with multiple outputs

5 files changed, 50 insertions, 6 deletions

diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowImporter.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowImporter.java
index 7116d430502..9ff88103f12 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowImporter.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowImporter.java
@@ -82,8 +82,6 @@ public class TensorFlowImporter {
         findDimensionNames(model, index);
         importExpressions(model, index, bundle);
 
-        // nodes with multiple outputs are calculated multiple times. consider adding macros for those.
-
         reportWarnings(model, index);
 
         return model;
@@ -241,7 +239,14 @@ public class TensorFlowImporter {
 
     private static void importMacroExpression(TensorFlowModel model, TensorFlowOperation operation) {
         if (operation.macro().isPresent()) {
-            model.macro(operation.vespaName(), operation.macro().get());
+            TensorFunction function = operation.macro().get();
+            try {
+                model.macro(operation.macroName(), new RankingExpression(operation.macroName(), function.toString()));
+            }
+            catch (ParseException e) {
+                throw new RuntimeException("Tensorflow function " + function +
+                        " cannot be parsed as a ranking expression", e);
+            }
         }
     }
 
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/PlaceholderWithDefault.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/PlaceholderWithDefault.java
index ab091b77a65..4e5709505ce 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/PlaceholderWithDefault.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/PlaceholderWithDefault.java
@@ -34,7 +34,7 @@ public class PlaceholderWithDefault extends TensorFlowOperation {
     }
 
     @Override
-    public Optional<RankingExpression> macro() {
+    public Optional<TensorFunction> macro() {
         // For now, it is much more efficient to assume we always will return
         // the default value, as we can prune away large parts of the expression
         // tree by having it calculated as a constant. If a case arises where
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/TensorFlowOperation.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/TensorFlowOperation.java
index fd9dfd167fb..9e8f6df3e2c 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/TensorFlowOperation.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/TensorFlowOperation.java
@@ -8,6 +8,7 @@ import com.yahoo.searchlib.rankingexpression.integration.tensorflow.importer.Ord
 import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
 import com.yahoo.searchlib.rankingexpression.rule.ReferenceNode;
 import com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode;
+import com.yahoo.tensor.evaluation.VariableTensor;
 import com.yahoo.tensor.functions.TensorFunction;
 import org.tensorflow.framework.NodeDef;
 
@@ -28,6 +29,8 @@ import java.util.function.Function;
  */
 public abstract class TensorFlowOperation {
 
+    protected final static String MACRO_PREFIX = "tf_macro_";
+
     protected final NodeDef node;
     protected final int port;
     protected final List<TensorFlowOperation> inputs;
@@ -36,6 +39,7 @@ public abstract class TensorFlowOperation {
 
     protected OrderedTensorType type;
     protected TensorFunction function;
+    protected TensorFunction macro = null;
 
     private Value constantValue = null;
     private List<TensorFlowOperation> controlInputs = Collections.emptyList();
@@ -65,6 +69,9 @@ public abstract class TensorFlowOperation {
             if (isConstant()) {
                 ExpressionNode constant = new ReferenceNode("constant(\"" + vespaName() + "\")");
                 function = new TensorFunctionNode.TensorFunctionExpressionNode(constant);
+            } else if (outputs.size() > 1) {
+                macro = lazyGetFunction();
+                function = new VariableTensor(macroName(), type.type());
             } else {
                 function = lazyGetFunction();
             }
@@ -82,7 +89,7 @@ public abstract class TensorFlowOperation {
     public List<TensorFlowOperation> outputs() { return Collections.unmodifiableList(outputs); }
 
     /** Returns a Vespa ranking expression that should be added as a macro */
-    public Optional<RankingExpression> macro() { return Optional.empty(); }
+    public Optional<TensorFunction> macro() { return Optional.ofNullable(macro); }
 
     /** Add dimension name constraints for this operation */
     public void addDimensionNameConstraints(DimensionRenamer renamer) { }
@@ -111,6 +118,9 @@ public abstract class TensorFlowOperation {
     /** Retrieve the valid Vespa name of this node */
     public String vespaName() { return node.getName() != null ? node.getName().replace('/', '_') : null; }
 
+    /** Retrieve the valid Vespa name of this node if it is a macro */
+    public String macroName() { return vespaName() != null ? MACRO_PREFIX + vespaName() : null; }
+
     /** Retrieve the list of warnings produced during its lifetime */
     public List<String> warnings() { return Collections.unmodifiableList(importWarnings); }
 
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/DropoutImportTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/DropoutImportTestCase.java
index f64d697d9b9..c09b1f2b606 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/DropoutImportTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/DropoutImportTestCase.java
@@ -32,7 +32,7 @@ public class DropoutImportTestCase {
         RankingExpression output = signature.outputExpression("y");
         assertNotNull(output);
         assertEquals("outputs/BiasAdd", output.getName());
-        assertEquals("join(reduce(join(rename(X, (d0, d1), (d0, d2)), constant(\"outputs_kernel_read\"), f(a,b)(a * b)), sum, d2), constant(\"outputs_bias_read\"), f(a,b)(a + b))",
+        assertEquals("join(reduce(join(tf_macro_X, constant(\"outputs_kernel_read\"), f(a,b)(a * b)), sum, d2), constant(\"outputs_bias_read\"), f(a,b)(a + b))",
                 output.getRoot().toString());
         model.assertEqualResult("X", output.getName());
     }
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 1691756a64d..9f372d8d6f5 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
@@ -1,10 +1,14 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.searchlib.rankingexpression.integration.tensorflow;
 
+import com.yahoo.searchlib.rankingexpression.RankingExpression;
 import com.yahoo.searchlib.rankingexpression.evaluation.Context;
 import com.yahoo.searchlib.rankingexpression.evaluation.MapContext;
 import com.yahoo.searchlib.rankingexpression.evaluation.TensorValue;
 import com.yahoo.searchlib.rankingexpression.integration.tensorflow.importer.TensorConverter;
+import com.yahoo.searchlib.rankingexpression.rule.CompositeNode;
+import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
+import com.yahoo.searchlib.rankingexpression.rule.ReferenceNode;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorType;
 import org.tensorflow.SavedModelBundle;
@@ -42,6 +46,9 @@ public class TestableTensorFlowModel {
         Context context = contextFrom(model);
         Tensor placeholder = placeholderArgument();
         context.put(inputName, new TensorValue(placeholder));
+
+        model.macros().forEach((k,v) -> evaluateMacro(context, model, k));
+
         Tensor vespaResult = model.expressions().get(operationName).evaluate(context).asTensor();
         assertEquals("Operation '" + operationName + "' produces equal results", tfResult, vespaResult);
     }
@@ -74,4 +81,26 @@ public class TestableTensorFlowModel {
         return b.build();
     }
 
+    private void evaluateMacro(Context context, TensorFlowModel model, String macroName) {
+        if (!context.names().contains(macroName)) {
+            RankingExpression e = model.macros().get(macroName);
+            evaluateMacroDependencies(context, model, e.getRoot());
+            context.put(macroName, new TensorValue(e.evaluate(context).asTensor()));
+        }
+    }
+
+    private void evaluateMacroDependencies(Context context, TensorFlowModel model, ExpressionNode node) {
+        if (node instanceof ReferenceNode) {
+            String name = node.toString();
+            if (model.macros().containsKey(name)) {
+                evaluateMacro(context, model, name);
+            }
+        }
+        else if (node instanceof CompositeNode) {
+            for (ExpressionNode child : ((CompositeNode)node).children()) {
+                evaluateMacroDependencies(context, model, child);
+            }
+        }
+    }
+
 }