Output intermediate graph with type info on error

8 files changed, 17 insertions, 18 deletions

diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ModelImporter.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ModelImporter.java
index b88d7423a82..f22e89cc8bb 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ModelImporter.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ModelImporter.java
@@ -11,8 +11,7 @@ import com.yahoo.searchlib.rankingexpression.parser.ParseException;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.functions.Rename;
 import com.yahoo.tensor.functions.TensorFunction;
-import com.yahoo.text.ParenthesisExpressionPrettyPrinter;
-import com.yahoo.text.Text;
+import com.yahoo.text.ExpressionFormatter;
 import com.yahoo.yolean.Exceptions;
 
 import java.io.File;
@@ -54,7 +53,7 @@ public abstract class ModelImporter implements MlModelImporter {
     protected static ImportedModel convertIntermediateGraphToModel(IntermediateGraph graph, String modelSource) {
         ImportedModel model = new ImportedModel(graph.name(), modelSource);
         log.log(Level.FINE, () -> "Intermediate graph created from '" + modelSource + "':\n" +
-                                  ParenthesisExpressionPrettyPrinter.prettyPrint(graph.toFullString()));
+                                  ExpressionFormatter.inTwoColumnMode(20).format(graph.toFullString()));
         System.out.println("Intermediate graph created from '" + modelSource + "':\n" + graph.toFullString());
 
         graph.optimize();
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java
index 7d1b6a61e2e..d5671889e01 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java
@@ -56,6 +56,7 @@ public class Argument extends IntermediateOperation {
 
     @Override
     public String toString() {
-        return "Argument(" + standardNamingType + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tArgument(" + standardNamingType + ")";
     }
+
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java
index 6571e77a198..e7c4b355217 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java
@@ -89,7 +89,7 @@ public class Const extends IntermediateOperation {
 
     @Override
     public String toString() {
-        return "Const(" + type + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tConst(" + type + ")";
     }
 
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/ExpandDims.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/ExpandDims.java
index 8f7d3755005..18ff602148a 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/ExpandDims.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/ExpandDims.java
@@ -30,7 +30,7 @@ public class ExpandDims extends IntermediateOperation {
         if ( ! allInputTypesPresent(2)) return null;
 
         IntermediateOperation axisOperation = inputs().get(1);
-        if (!axisOperation.getConstantValue().isPresent()) {
+        if (  !axisOperation.getConstantValue().isPresent()) {
             throw new IllegalArgumentException("ExpandDims in " + name + ": Axis must be a constant.");
         }
         Tensor axis = axisOperation.getConstantValue().get().asTensor();
@@ -103,8 +103,8 @@ public class ExpandDims extends IntermediateOperation {
 
     @Override
     public String toFullString() {
-        return "ExpandDims(" + inputs().get(0).toFullString() + ", " +
-               inputs().get(1).toFullString() + ", " + expandDimensions + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tExpandDims(" + inputs().get(0).toFullString() + ", " +
+               inputs().get(1).toFullString() + ", " + expandDimensions + ")";
     }
 
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Join.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Join.java
index c5e6ae49a25..5c7acc8a0ee 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Join.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Join.java
@@ -118,8 +118,8 @@ public class Join extends IntermediateOperation {
 
     @Override
     public String toFullString() {
-        return "Join(" + inputs().get(0).toFullString() + ", " +
-               inputs().get(1).toFullString() + ", " + operator + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tJoin(" + inputs().get(0).toFullString() + ", " +
+               inputs().get(1).toFullString() + ", " + operator + ")";
     }
 
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Map.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Map.java
index 4f70c46e459..5fba7cc6da0 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Map.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Map.java
@@ -41,7 +41,7 @@ public class Map extends IntermediateOperation {
 
     @Override
     public String toFullString() {
-        return "Map(" + inputs().get(0).toFullString() + ", " + operator + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + "\tMap(" + inputs().get(0).toFullString() + ", " + operator + ")";
     }
 
 }
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/MatMul.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/MatMul.java
index 73aa40927be..3f7c32294ed 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/MatMul.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/MatMul.java
@@ -5,8 +5,7 @@ import ai.vespa.rankingexpression.importer.DimensionRenamer;
 import ai.vespa.rankingexpression.importer.OrderedTensorType;
 import com.yahoo.tensor.TensorType;
 import com.yahoo.tensor.functions.TensorFunction;
-import com.yahoo.text.ParenthesisExpressionPrettyPrinter;
-import com.yahoo.text.Text;
+import com.yahoo.text.ExpressionFormatter;
 
 import java.util.List;
 import java.util.Optional;
@@ -81,13 +80,12 @@ public class MatMul extends IntermediateOperation {
 
         throw new IllegalArgumentException("Expected 2 dimensions in the " + inputDescription + " to " + this +
                                            " but got just " + dimensions + " from\n" +
-                                           ParenthesisExpressionPrettyPrinter.prettyPrint(supplier.toFullString()));
+                                           ExpressionFormatter.inTwoColumnMode(70).format(supplier.toFullString()));
     }
 
     @Override
     public String toFullString() {
-        return "MatMul(" + inputs().get(0).toFullString() + ", " +
-                           inputs().get(1).toFullString() + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tMatMul(" + inputs().get(0).toFullString() + ", " + inputs().get(1).toFullString() + ")";
     }
 
     @Override
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Sum.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Sum.java
index df5c4e9cbfa..8b5d06e8d56 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Sum.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Sum.java
@@ -104,6 +104,7 @@ public class Sum extends IntermediateOperation {
                 builder.add(TensorType.Dimension.indexed(dimension.name(), 1L));
             }
         }
+        System.out.println("----------> Sum input type is " + inputType + ", keepDimensions: " + keepDimensions + ", result: " + builder.build());
         return builder.build();
     }
 
@@ -114,8 +115,8 @@ public class Sum extends IntermediateOperation {
 
     @Override
     public String toFullString() {
-        return "Sum(" + inputs().get(0).toFullString() + ", " +
-               inputs().get(1).toFullString() + ", " + reduceDimensions + ")" + " : " + lazyGetType();
+        return "\t" + lazyGetType() + ":\tSum[keep_dims=" + shouldKeepDimensions() + "](" +
+               inputs().get(0).toFullString() + ", " + inputs().get(1).toFullString() + ", " + reduceDimensions + ")";
     }
 
 }