Output bound dense tensors in dense string form

5 files changed, 88 insertions, 19 deletions

diff --git a/config-model/src/test/java/com/yahoo/searchdefinition/processing/VespaMlModelTestCase.java b/config-model/src/test/java/com/yahoo/searchdefinition/processing/VespaMlModelTestCase.java
index 34b727f9f4e..412264aec57 100644
--- a/config-model/src/test/java/com/yahoo/searchdefinition/processing/VespaMlModelTestCase.java
+++ b/config-model/src/test/java/com/yahoo/searchdefinition/processing/VespaMlModelTestCase.java
@@ -26,7 +26,7 @@ public class VespaMlModelTestCase {
 
     private final String expectedRankConfig =
             "constant(constant1).type : tensor(x[3])\n" +
-            "constant(constant1).value : tensor(x[3]):{{x:0}:0.5,{x:1}:1.5,{x:2}:2.5}\n" +
+            "constant(constant1).value : tensor(x[3]):[0.5, 1.5, 2.5]\n" +
             "rankingExpression(foo1).rankingScript : reduce(reduce(input1 * input2, sum, name) * constant(constant1), max, x) * 3.0\n" +
             "rankingExpression(foo1).input2.type : tensor(x[3])\n" +
             "rankingExpression(foo1).input1.type : tensor(name{},x[3])\n" +
diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/vespa/VespaImportTestCase.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/vespa/VespaImportTestCase.java
index c7210e6710a..faa603d0ab0 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/vespa/VespaImportTestCase.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/vespa/VespaImportTestCase.java
@@ -31,11 +31,11 @@ public class VespaImportTestCase {
         assertEquals("tensor(x[3])", model.inputs().get("input2").toString());
 
         assertEquals(2, model.smallConstants().size());
-        assertEquals("tensor(x[3]):{{x:0}:0.5,{x:1}:1.5,{x:2}:2.5}", model.smallConstants().get("constant1"));
+        assertEquals("tensor(x[3]):[0.5, 1.5, 2.5]", model.smallConstants().get("constant1"));
         assertEquals("tensor():{3.0}", model.smallConstants().get("constant2"));
 
         assertEquals(1, model.largeConstants().size());
-        assertEquals("tensor(x[3]):{{x:0}:0.5,{x:1}:1.5,{x:2}:2.5}", model.largeConstants().get("constant1asLarge"));
+        assertEquals("tensor(x[3]):[0.5, 1.5, 2.5]", model.largeConstants().get("constant1asLarge"));
 
         assertEquals(2, model.expressions().size());
         assertEquals("reduce(reduce(input1 * input2, sum, name) * constant1, max, x) * constant2",
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
index aca2bfc1b0f..a3c30af66ec 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
@@ -210,7 +210,36 @@ public abstract class IndexedTensor implements Tensor {
     }
 
     @Override
-    public String toString() { return Tensor.toStandardString(this); }
+    public String toString() {
+        if (type.rank() == 0) return Tensor.toStandardString(this);
+        if (type.dimensions().stream().anyMatch(d -> d.size().isEmpty())) return Tensor.toStandardString(this);
+
+        Indexes indexes = Indexes.of(dimensionSizes);
+
+        StringBuilder b = new StringBuilder(type.toString()).append(":");
+        for (int index = 0; index < size(); index++) {
+            indexes.next();
+
+            // start brackets
+            for (int i = 0; i < indexes.rightDimensionsWhichAreAtStart(); i++)
+                b.append("[");
+
+            // value
+            if (type.valueType() == TensorType.Value.DOUBLE)
+                b.append(get(index));
+            else if (type.valueType() == TensorType.Value.FLOAT)
+                b.append(getFloat(index));
+            else
+                throw new IllegalStateException("Unexpected value type " + type.valueType());
+
+            // end bracket and comma
+            for (int i = 0; i < indexes.rightDimensionsWhichAreAtEnd(); i++)
+                b.append("]");
+            if (index < size() - 1)
+                b.append(", ");
+        }
+        return b.toString();
+    }
 
     @Override
     public boolean equals(Object other) {
@@ -827,6 +856,27 @@ public abstract class IndexedTensor implements Tensor {
 
         public abstract void next();
 
+        /** Returns the number of dimensions from the right which are currently at the start position (0) */
+        int rightDimensionsWhichAreAtStart() {
+            int dimension = indexes.length - 1;
+            int atStartCount = 0;
+            while (dimension >= 0 && indexes[dimension] == 0) {
+                atStartCount++;
+                dimension--;
+            }
+            return atStartCount;
+        }
+
+        /** Returns the number of dimensions from the right which are currently at the end position */
+        int rightDimensionsWhichAreAtEnd() {
+            int dimension = indexes.length - 1;
+            int atEndCount = 0;
+            while (dimension >= 0 && indexes[dimension] == dimensionSizes().size(dimension) - 1) {
+                atEndCount++;
+                dimension--;
+            }
+            return atEndCount;
+        }
     }
 
     private final static class EmptyIndexes extends Indexes {
diff --git a/vespajlib/src/test/java/com/yahoo/tensor/TensorParserTestCase.java b/vespajlib/src/test/java/com/yahoo/tensor/TensorParserTestCase.java
index 63fe40565bd..1928971820c 100644
--- a/vespajlib/src/test/java/com/yahoo/tensor/TensorParserTestCase.java
+++ b/vespajlib/src/test/java/com/yahoo/tensor/TensorParserTestCase.java
@@ -23,37 +23,42 @@ public class TensorParserTestCase {
 
     @Test
     public void testDenseParsing() {
-        assertEquals(Tensor.Builder.of(TensorType.fromSpec("tensor()")).build(),
-                     Tensor.from("tensor():[]"));
-        assertEquals(Tensor.Builder.of(TensorType.fromSpec("tensor(x[1])")).cell(1.0, 0).build(),
-                     Tensor.from("tensor(x[1]):[1.0]"));
-        assertEquals(Tensor.Builder.of(TensorType.fromSpec("tensor(x[2])")).cell(1.0, 0).cell(2.0, 1).build(),
-                     Tensor.from("tensor(x[2]):[1.0, 2.0]"));
-        assertEquals(Tensor.Builder.of(TensorType.fromSpec("tensor(x[2],y[3])"))
+        assertDense(Tensor.Builder.of(TensorType.fromSpec("tensor()")).build(),
+                    "tensor():{0.0}");
+        assertDense(Tensor.Builder.of(TensorType.fromSpec("tensor()")).cell(1.3).build(),
+                    "tensor():{1.3}");
+        assertDense(Tensor.Builder.of(TensorType.fromSpec("tensor(x[])")).cell(1.0, 0).build(),
+                    "tensor(x[]):{{x:0}:1.0}");
+        assertDense(Tensor.Builder.of(TensorType.fromSpec("tensor(x[1])")).cell(1.0, 0).build(),
+                    "tensor(x[1]):[1.0]");
+        assertDense(Tensor.Builder.of(TensorType.fromSpec("tensor(x[2])")).cell(1.0, 0).cell(2.0, 1).build(),
+                    "tensor(x[2]):[1.0, 2.0]");
+        assertDense(Tensor.Builder.of(TensorType.fromSpec("tensor(x[2],y[3])"))
                                    .cell(1.0, 0, 0)
                                    .cell(2.0, 0, 1)
                                    .cell(3.0, 0, 2)
                                    .cell(4.0, 1, 0)
                                    .cell(5.0, 1, 1)
                                    .cell(6.0, 1, 2).build(),
-                     Tensor.from("tensor(x[2],y[3]):[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]"));
-        assertEquals(Tensor.Builder.of(TensorType.fromSpec("tensor(x[1],y[2],z[3])"))
+                    "tensor(x[2],y[3]):[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]");
+        assertDense(Tensor.Builder.of(TensorType.fromSpec("tensor(x[1],y[2],z[3])"))
                                    .cell(1.0, 0, 0, 0)
                                    .cell(2.0, 0, 0, 1)
                                    .cell(3.0, 0, 0, 2)
                                    .cell(4.0, 0, 1, 0)
                                    .cell(5.0, 0, 1, 1)
                                    .cell(6.0, 0, 1, 2).build(),
-                     Tensor.from("tensor(x[1],y[2],z[3]):[[[1.0], [2.0]], [[3.0], [4.0]], [[5.0], [6.0]]]"));
-        assertEquals(Tensor.Builder.of(TensorType.fromSpec("tensor(x[3],y[2],z[1])"))
+                    "tensor(x[1],y[2],z[3]):[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]");
+        assertDense(Tensor.Builder.of(TensorType.fromSpec("tensor(x[3],y[2],z[1])"))
                                    .cell(1.0, 0, 0, 0)
                                    .cell(2.0, 0, 1, 0)
                                    .cell(3.0, 1, 0, 0)
                                    .cell(4.0, 1, 1, 0)
                                    .cell(5.0, 2, 0, 0)
                                    .cell(6.0, 2, 1, 0).build(),
-                     Tensor.from("tensor(x[3],y[2],z[1]):[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]"));
-        assertEquals(Tensor.Builder.of(TensorType.fromSpec("tensor(x[3],y[2],z[1])"))
+                    "tensor(x[3],y[2],z[1]):[[[1.0], [2.0]], [[3.0], [4.0]], [[5.0], [6.0]]]");
+        assertEquals("Messy input",
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[3],y[2],z[1])"))
                                    .cell( 1.0, 0, 0, 0)
                                    .cell( 2.0, 0, 1, 0)
                                    .cell( 3.0, 1, 0, 0)
@@ -61,6 +66,20 @@ public class TensorParserTestCase {
                                    .cell( 5.0, 2, 0, 0)
                                    .cell(-6.0, 2, 1, 0).build(),
                      Tensor.from("tensor( x[3],y[2],z[1]) : [  [ [1.0, 2.0, 3.0] , [4.0, 5,-6.0] ]  ]"));
+        assertEquals("Skipping syntactic sugar",
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[3],y[2],z[1])"))
+                                   .cell( 1.0, 0, 0, 0)
+                                   .cell( 2.0, 0, 1, 0)
+                                   .cell( 3.0, 1, 0, 0)
+                                   .cell( 4.0, 1, 1, 0)
+                                   .cell( 5.0, 2, 0, 0)
+                                   .cell(-6.0, 2, 1, 0).build(),
+                     Tensor.from("tensor( x[3],y[2],z[1]) : [1.0, 2.0, 3.0 , 4.0, 5, -6.0]"));
+    }
+
+    private void assertDense(Tensor expectedTensor, String denseFormat) {
+        assertEquals(denseFormat, expectedTensor, Tensor.from(denseFormat));
+        assertEquals(denseFormat, expectedTensor.toString());
     }
 
     @Test
diff --git a/vespajlib/src/test/java/com/yahoo/tensor/TensorTestCase.java b/vespajlib/src/test/java/com/yahoo/tensor/TensorTestCase.java
index c53db160806..3d5d8d1f5ae 100644
--- a/vespajlib/src/test/java/com/yahoo/tensor/TensorTestCase.java
+++ b/vespajlib/src/test/java/com/yahoo/tensor/TensorTestCase.java
@@ -89,9 +89,9 @@ public class TensorTestCase {
 
     @Test
     public void testCombineInDimensionIndexed() {
-        Tensor input =  Tensor.from("tensor(input[]):{{input:0}:3, {input:1}:7}");
+        Tensor input =  Tensor.from("tensor(input[2]):{{input:0}:3, {input:1}:7}");
         Tensor result = input.concat(11, "input");
-        assertEquals("tensor(input[]):{{input:0}:3.0,{input:1}:7.0,{input:2}:11.0}", result.toString());
+        assertEquals("tensor(input[3]):[3.0, 7.0, 11.0]", result.toString());
     }
 
     /** All functions are more throughly tested in searchlib EvaluationTestCase */