More functions. Cleanup.

2 files changed, 53 insertions, 39 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 8dcd31b270e..e47f2ad53d9 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
@@ -7,10 +7,8 @@ import com.yahoo.io.IOUtils;
 import com.yahoo.searchlib.rankingexpression.RankingExpression;
 import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
 import com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode;
-import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorType;
 import com.yahoo.tensor.evaluation.VariableTensor;
-import com.yahoo.tensor.functions.ConstantTensor;
 import com.yahoo.tensor.functions.Join;
 import com.yahoo.tensor.functions.Matmul;
 import com.yahoo.tensor.functions.Rename;
@@ -34,6 +32,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.DoubleBinaryOperator;
+import java.util.function.DoubleUnaryOperator;
 import java.util.stream.Collectors;
 
 /**
@@ -146,11 +145,13 @@ public class TensorFlowImporter {
                                                  GraphDef graph,
                                                  String indent) {
         // Import arguments lazily below, as some nodes have arguments unused arguments leading to unsupported ops
-        switch (tfNode.getOp()) {
-            case "Identity" : return identity(tfNode, inputs);
-            case "Add" : return join(importArguments(tfNode, inputs, graph, indent), ScalarFunctions.add());
-            case "MatMul" : return matmul(importArguments(tfNode, inputs, graph, indent));
-            case "Softmax" : return softmax(importArguments(tfNode, inputs, graph, indent));
+        // TODO: Implement mapping of more functions from https://www.tensorflow.org/api_docs/python/
+        switch (tfNode.getOp().toLowerCase()) {
+            case "add" : case "add_n" : return join(importArguments(tfNode, inputs, graph, indent), ScalarFunctions.add());
+            case "acos" : return map(importArguments(tfNode, inputs, graph, indent), ScalarFunctions.acos());
+            case "identity" : return identity(tfNode, inputs);
+            case "matmul" : return matmul(importArguments(tfNode, inputs, graph, indent));
+            case "softmax" : return softmax(importArguments(tfNode, inputs, graph, indent));
             default : throw new IllegalArgumentException("Conversion of TensorFlow operation '" + tfNode.getOp() + "' is not supported");
         }
     }
@@ -162,15 +163,51 @@ public class TensorFlowImporter {
     }
     
     private TypedTensorFunction join(List<TypedTensorFunction> arguments, DoubleBinaryOperator doubleFunction) {
+        // Note that this generalizes the corresponding TF function as it does not verify that the tensor
+        // types are the same, with the assumption that this already happened on the TF side
+        // (and if not, this should do the right thing anyway)
         ensureArguments(2, arguments, "join");
         TypedTensorFunction a = arguments.get(0);
         TypedTensorFunction b = arguments.get(0);
-        // TODO: Verify with TF doc
-        TensorType resultType = Join.resultType(a.type(), b.type());
+
+        TensorType resultType = Join.outputType(a.type(), b.type());
         Join function = new Join(a.function(), b.function(), doubleFunction);
         return new TypedTensorFunction(resultType, function);
     }
-    
+
+    private TypedTensorFunction map(List<TypedTensorFunction> arguments, DoubleUnaryOperator doubleFunction) {
+        ensureArguments(1, arguments, "apply");
+        TypedTensorFunction a = arguments.get(0);
+
+        TensorType resultType = com.yahoo.tensor.functions.Map.outputType(a.type());
+        com.yahoo.tensor.functions.Map function = new com.yahoo.tensor.functions.Map(a.function(), doubleFunction);
+        return new TypedTensorFunction(resultType, function);
+    }
+
+    private TypedTensorFunction identity(NodeDef tfNode, Map<String, TensorType> inputs) {
+        // TODO: Verify with TF documentation
+        String name;
+        TensorType inputType;
+        if (tfNode.getName().endsWith("/read")) { // A node reading a variable supplied with this model TODO: We need to turn those into constants
+            if (tfNode.getInputList().size() != 1)
+                throw new IllegalArgumentException("A Variable/read node must have one input but has " +
+                                                   tfNode.getInputList().size());
+            name = tfNode.getInput(0);
+            AttrValue shapes = tfNode.getAttrMap().get("_output_shapes");
+            if (shapes == null)
+                throw new IllegalArgumentException("Referenced variable '" + name + " is missing a tensor output shape");
+            inputType = importTensorType(shapes.getList().getShape(0));
+        }
+        else { // a referenced input (query or document tensor) TODO: How to map to attribute/query name
+            name = tfNode.getName();
+            inputType = inputs.get(name);
+            if (inputType == null)
+                throw new IllegalArgumentException("An identity operation node is referencing input '" + name +
+                                                   "', but there is no such input");
+        }
+        return new TypedTensorFunction(inputType, new VariableTensor(name));
+    }
+
     private TypedTensorFunction matmul(List<TypedTensorFunction> arguments) {
         ensureArguments(2, arguments, "matmul");
         TypedTensorFunction a = arguments.get(0);
@@ -183,7 +220,7 @@ public class TensorFlowImporter {
         // Let the second-to-last dimension of the second tensor be the same as the last dimension of the first
         // and the last dimension of the second argument be not present in the first argument, while leaving the 
         // rest of the dimensions the same. Such is the way of implicit dimension name tensor multiplication.
-        
+
         // TODO: Check if transpose_a or transpose_b is set and rename differently accordingly
 
         String beforeLastDim = "d" + (a.type().rank() - 1);
@@ -193,40 +230,17 @@ public class TensorFlowImporter {
         Rename renamedB = new Rename(b.function(), ImmutableList.of(beforeLastDim, lastDim), 
                                                    ImmutableList.of(lastDim, afterLastDim));
         Matmul matmul = new Matmul(a.function(), renamedB, lastDim);
-        return new TypedTensorFunction(Matmul.resultType(a.type(), b.type(), lastDim), 
+        return new TypedTensorFunction(Matmul.outputType(a.type(), b.type(), lastDim), 
                                        new Rename(matmul, afterLastDim, lastDim));
     }
 
     private TypedTensorFunction softmax(List<TypedTensorFunction> arguments) {
         ensureArguments(1, arguments, "softmax");
         TypedTensorFunction a = arguments.get(0);
-        String dimension = "d0";  // TODO: Verify with TF doc
+        // TODO: Read the "dim" parameter and use it to decide dimension if set and != -1
+        String dimension = "d" + (a.type().rank() - 1);
         Softmax softmax = new Softmax(a.function(), dimension);
-        return new TypedTensorFunction(Softmax.resultType(a.type(), dimension), softmax);
-    }
-
-    private TypedTensorFunction identity(NodeDef tfNode, Map<String, TensorType> inputs) {
-        // TODO: Verify with TF documentation
-        String name;
-        TensorType inputType;
-        if (tfNode.getName().endsWith("/read")) { // A node reading a variable supplied with this model TODO: We need to turn those into constants
-            if (tfNode.getInputList().size() != 1)
-                throw new IllegalArgumentException("A Variable/read node must have one input but has " +
-                                                   tfNode.getInputList().size());
-            name = tfNode.getInput(0);
-            AttrValue shapes = tfNode.getAttrMap().get("_output_shapes");
-            if (shapes == null)
-                throw new IllegalArgumentException("Referenced variable '" + name + " is missing a tensor output shape");
-            inputType = importTensorType(shapes.getList().getShape(0));
-        }
-        else { // a referenced input (query or document tensor) TODO: How to map to attribute/query name
-            name = tfNode.getName();
-            inputType = inputs.get(name);
-            if (inputType == null)
-                throw new IllegalArgumentException("An identity operation node is referencing input '" + name +
-                                                   "', but there is no such input");
-        }
-        return new TypedTensorFunction(inputType, new VariableTensor(name));
+        return new TypedTensorFunction(Softmax.outputType(a.type(), dimension), softmax);
     }
 
     private void ensureArguments(int count, List<TypedTensorFunction> arguments, String operationName) {
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowImporterTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowImporterTestCase.java
index 30328c3d9fe..bfe2bb3a63b 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowImporterTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/TensorFlowImporterTestCase.java
@@ -23,7 +23,7 @@ public class TensorFlowImporterTestCase {
                      "softmax(join(rename(matmul(x, rename(x, (d1, d2), (d2, d3)), d2), d3, d2), " +
                                   "rename(matmul(x, rename(x, (d1, d2), (d2, d3)), d2), d3, d2), " +
                                   "f(a,b)(a + b)), " +
-                             "d0)",
+                             "d1)",
                      toNonPrimitiveString(expressions.get(0)));
     }