Parse unbound tensors in short form

2 files changed, 125 insertions, 5 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/TensorParser.java b/vespajlib/src/main/java/com/yahoo/tensor/TensorParser.java
index 0a1d9b6cf6e..eaa6e50e87f 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/TensorParser.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/TensorParser.java
@@ -124,15 +124,30 @@ class TensorParser {
         if (type.isEmpty())
             throw new IllegalArgumentException("The dense tensor form requires an explicit tensor type " +
                                                "on the form 'tensor(dimensions):...");
-        if (type.get().dimensions().stream().anyMatch(d -> (d.size().isEmpty())))
-            throw new IllegalArgumentException("The dense tensor form requires a tensor type containing " +
-                                               "only dense dimensions with a given size");
 
-        IndexedTensor.BoundBuilder builder = (IndexedTensor.BoundBuilder) IndexedTensor.Builder.of(type.get());
-        new DenseValueParser(valueString, dimensionOrder, builder).parse();
+        IndexedTensor.Builder builder = IndexedTensor.Builder.of(type.get());
+
+        if (type.get().dimensions().stream().anyMatch(d -> (d.size().isEmpty()))) {
+            new UnboundDenseValueParser(valueString, builder).parse();
+            return checkBoundDimensionSizes(builder.build());
+        }
+
+        new DenseValueParser(valueString, dimensionOrder, (IndexedTensor.BoundBuilder) builder).parse();
         return builder.build();
     }
 
+    private static Tensor checkBoundDimensionSizes(IndexedTensor tensor) {
+        TensorType type = tensor.type();
+        for (int i = 0; i < type.dimensions().size(); ++i) {
+            TensorType.Dimension dimension = type.dimensions().get(i);
+            if (dimension.size().isPresent() && dimension.size().get() != tensor.dimensionSizes().size(i)) {
+                throw new IllegalArgumentException("Unexpected size " + tensor.dimensionSizes().size(i) +
+                        " for dimension " + dimension.name() + " for type " + type);
+            }
+        }
+        return tensor;
+    }
+
     private static abstract class ValueParser {
 
         protected final String string;
@@ -299,6 +314,64 @@ class TensorParser {
     }
 
     /**
+     * Parses unbound tensor short forms - e.g. tensor(x[],y[]):[[1,2,3],[4,5,6]]
+     */
+    private static class UnboundDenseValueParser extends ValueParser {
+
+        private final IndexedTensor.Builder builder;
+        private final long[] indexes;
+
+        public UnboundDenseValueParser(String string, IndexedTensor.Builder builder) {
+            super(string);
+            this.builder = builder;
+            this.indexes = new long[builder.type().dimensions().size()];
+        }
+
+        public void parse() {
+            consumeList(0);
+        }
+
+        private void consumeList(int dimension) {
+            consume('[');
+            indexes[dimension] = 0;
+            while ( ! atListEnd() ) {
+                if (isInnerMostDimension(dimension)) {
+                    consumeNumber();
+                } else {
+                    consumeList(dimension + 1);
+                }
+                indexes[dimension]++;
+                consumeOptional(',');
+            }
+            consume(']');
+        }
+
+        private void consumeNumber() {
+            Number number = consumeNumber(builder.type().valueType());
+            switch (builder.type().valueType()) {
+                case DOUBLE:   builder.cell((Double)number, indexes); break;
+                case FLOAT:    builder.cell((Float)number, indexes); break;
+                case BFLOAT16: builder.cell((Float)number, indexes); break;
+                case INT8:     builder.cell((Float)number, indexes); break;
+            }
+        }
+
+        private boolean isInnerMostDimension(int dimension) {
+            return dimension == (indexes.length - 1);
+        }
+
+        protected boolean atListEnd() {
+            skipSpace();
+            if (position >= string.length()) {
+                throw new IllegalArgumentException("At value position " + position + ": Expected a ']'" +
+                        " but got the end of the string");
+            }
+            return string.charAt(position) == ']';
+        }
+
+    }
+
+    /**
      * Parses mixed tensor short forms {a:[1,2], ...} AND 1d mapped tensor short form {a:b, ...}.
      */
     private static class MixedValueParser extends ValueParser {
diff --git a/vespajlib/src/test/java/com/yahoo/tensor/TensorParserTestCase.java b/vespajlib/src/test/java/com/yahoo/tensor/TensorParserTestCase.java
index 431e4b06263..b869107e744 100644
--- a/vespajlib/src/test/java/com/yahoo/tensor/TensorParserTestCase.java
+++ b/vespajlib/src/test/java/com/yahoo/tensor/TensorParserTestCase.java
@@ -150,6 +150,53 @@ public class TensorParserTestCase {
                      Tensor.from("tensor(key{},y[2],x[3]):{key1:[[1,2,3],[4,5,6]], key2:[[7,8,9],[10,11,12]]}"));
     }
 
+    @Test
+    public void testUnboundShortFormParsing() {
+        assertEquals(Tensor.from("tensor(x[]):[1.0, 2.0]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[])")).cell(1.0, 0).cell(2.0, 1).build());
+        assertEquals(Tensor.from("tensor<float>(x[]):[1.0, 2.0]"),
+                Tensor.Builder.of(TensorType.fromSpec("tensor<float>(x[])")).cell(1.0, 0).cell(2.0, 1).build());
+        assertEquals(Tensor.from("tensor<int8>(x[]):[1.0, 2.0]"),
+                Tensor.Builder.of(TensorType.fromSpec("tensor<int8>(x[])")).cell(1.0, 0).cell(2.0, 1).build());
+        assertEquals(Tensor.from("tensor<bfloat16>(x[]):[1.0, 2.0]"),
+                Tensor.Builder.of(TensorType.fromSpec("tensor<bfloat16>(x[])")).cell(1.0, 0).cell(2.0, 1).build());
+
+        assertEquals(Tensor.from("tensor(x[],y[]):[[1,2,3,4]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[])"))
+                        .cell(1.0, 0, 0).cell(2.0, 0, 1).cell(3.0, 0, 2).cell(4.0, 0, 3).build());
+        assertEquals(Tensor.from("tensor(x[],y[]):[[1,2],[3,4]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[])"))
+                         .cell(1.0, 0, 0).cell(2.0, 0, 1).cell(3.0, 1, 0).cell(4.0, 1, 1).build());
+        assertEquals(Tensor.from("tensor(x[],y[]):[[1],[2],[3],[4]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[])"))
+                        .cell(1.0, 0, 0).cell(2.0, 1, 0).cell(3.0, 2, 0).cell(4.0, 3, 0).build());
+        assertEquals(Tensor.from("tensor(x[],y[],z[]):[[[1,2],[3,4]]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[],z[])"))
+                        .cell(1.0, 0, 0, 0).cell(2.0, 0, 0, 1).cell(3.0, 0, 1, 0).cell(4.0, 0, 1, 1).build());
+        assertEquals(Tensor.from("tensor(x[],y[],z[]):[[[1],[2],[3],[4]]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[],z[])"))
+                        .cell(1.0, 0, 0, 0).cell(2.0, 0, 1, 0).cell(3.0, 0, 2, 0).cell(4.0, 0, 3, 0).build());
+        assertEquals(Tensor.from("tensor(x[],y[],z[]):[[[1,2,3,4]]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[],z[])"))
+                        .cell(1.0, 0, 0, 0).cell(2.0, 0, 0, 1).cell(3.0, 0, 0, 2).cell(4.0, 0, 0, 3).build());
+        assertEquals(Tensor.from("tensor(x[],y[],z[]):[[[1]],[[2]],[[3]],[[4]]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[],z[])"))
+                        .cell(1.0, 0, 0, 0).cell(2.0, 1, 0, 0).cell(3.0, 2, 0, 0).cell(4.0, 3, 0, 0).build());
+        assertEquals(Tensor.from("tensor(x[],y[],z[]):[[[1, 2]],[[3, 4]]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[],z[])"))
+                        .cell(1.0, 0, 0, 0).cell(2.0, 0, 0, 1).cell(3.0, 1, 0, 0).cell(4.0, 1, 0, 1).build());
+
+        assertEquals(Tensor.from("tensor(x[],y[],z[4]):[[[1,2,3,4]]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[],z[])"))
+                        .cell(1.0, 0, 0, 0).cell(2.0, 0, 0, 1).cell(3.0, 0, 0, 2).cell(4.0, 0, 0, 3).build());
+        assertEquals(Tensor.from("tensor(x[2],y[],z[2]):[[[1, 2]],[[3, 4]]]"),
+                     Tensor.Builder.of(TensorType.fromSpec("tensor(x[],y[],z[])"))
+                        .cell(1.0, 0, 0, 0).cell(2.0, 0, 0, 1).cell(3.0, 1, 0, 0).cell(4.0, 1, 0, 1).build());
+
+        assertIllegal("Unexpected size 2 for dimension y for type tensor(x[],y[3])",
+                "tensor(x[],y[3]):[[1,2],[3,4]]");
+    }
+
     private void assertDense(Tensor expectedTensor, String denseFormat) {
         assertEquals(denseFormat, expectedTensor, Tensor.from(denseFormat));
         assertEquals(denseFormat, expectedTensor.toString());