tensor peek

4 files changed, 134 insertions, 12 deletions

diff --git a/eval/src/tests/eval/compiled_function/compiled_function_test.cpp b/eval/src/tests/eval/compiled_function/compiled_function_test.cpp
index 151a4cf5dd5..1d5a6929083 100644
--- a/eval/src/tests/eval/compiled_function/compiled_function_test.cpp
+++ b/eval/src/tests/eval/compiled_function/compiled_function_test.cpp
@@ -58,6 +58,9 @@ std::vector<vespalib::string> unsupported = {
 };
 
 bool is_unsupported(const vespalib::string &expression) {
+    if (expression.find("{") != vespalib::string::npos) {
+        return true;
+    }
     for (const auto &prefix: unsupported) {
         if (starts_with(expression, prefix)) {
             return true;
diff --git a/eval/src/tests/eval/function/function_test.cpp b/eval/src/tests/eval/function/function_test.cpp
index 0e3100ae425..d93b57ab5b2 100644
--- a/eval/src/tests/eval/function/function_test.cpp
+++ b/eval/src/tests/eval/function/function_test.cpp
@@ -81,6 +81,12 @@ void verify_error(const vespalib::string &expr, const vespalib::string &expected
     EXPECT_EQUAL(expected_error, function.get_error());
 }
 
+void verify_parse(const vespalib::string &expr, const vespalib::string &expect) {
+    Function function = Function::parse(expr);
+    EXPECT_TRUE(!function.has_error());
+    EXPECT_EQUAL(function.dump_as_lambda(), expect);
+}
+
 TEST("require that scientific numbers can be parsed") {
     EXPECT_EQUAL(1.0,     as_number(Function::parse(params, "1")));
     EXPECT_EQUAL(2.5,     as_number(Function::parse(params, "2.5")));
@@ -933,6 +939,28 @@ TEST("require that convenient tensor create detects under-specified cells") {
 
 //-----------------------------------------------------------------------------
 
+TEST("require that tensor peek can be parsed") {
+    TEST_DO(verify_parse("t{x:1,y:foo}", "f(t)(t{x:1,y:foo})"));
+}
+
+TEST("require that tensor peek can contain expressions") {
+    TEST_DO(verify_parse("t{x:1+2,y:foo}", "f(t)(t{x:(1+2),y:foo})"));
+    TEST_DO(verify_parse("t{x:1+bar,y:foo}", "f(t,bar)(t{x:(1+bar),y:foo})"));
+    TEST_DO(verify_parse("t{x:1,y:foo+2}", "f(t,foo)(t{x:1,y:(foo+2)})"));
+    TEST_DO(verify_parse("t{x:1,y:(foo)}", "f(t,foo)(t{x:1,y:(foo)})"));
+}
+
+TEST("require that tensor peek can contain extra whitespace") {
+    TEST_DO(verify_parse(" t { x : 1 + bar , y : foo + 2 } ",
+                         "f(t,bar,foo)(t{x:(1+bar),y:(foo+2)})"));
+}
+
+TEST("require that empty tensor peek is not allowed") {
+    TEST_DO(verify_error("x{}", "[x{}]...[empty peek spec]...[]"));
+}
+
+//-----------------------------------------------------------------------------
+
 TEST("require that tensor concat can be parsed") {
     EXPECT_EQUAL("concat(a,b,d)", Function::parse({"a", "b"}, "concat(a,b,d)").dump());
     EXPECT_EQUAL("concat(a,b,d)", Function::parse({"a", "b"}, " concat ( a , b , d ) ").dump());
diff --git a/eval/src/tests/eval/node_types/node_types_test.cpp b/eval/src/tests/eval/node_types/node_types_test.cpp
index b2ad107f2aa..97a80d9b4b3 100644
--- a/eval/src/tests/eval/node_types/node_types_test.cpp
+++ b/eval/src/tests/eval/node_types/node_types_test.cpp
@@ -214,6 +214,31 @@ TEST("require that tensor create resolves correct type") {
     TEST_DO(verify("tensor(x[3]):{{x:0}:double,{x:1}:error,{x:2}:double}", "error"));
 }
 
+TEST("require that tensor peek resolves correct type") {
+    TEST_DO(verify("tensor(x[3]){x:1}", "double"));
+    TEST_DO(verify("tensor(x[3]){x:double}", "error"));
+    TEST_DO(verify("tensor(x[3]){x:(double)}", "double"));
+    TEST_DO(verify("tensor(x[3]){x:3}", "error"));
+    TEST_DO(verify("tensor(x{}){x:1}", "double"));
+    TEST_DO(verify("tensor(x{}){x:foo}", "double"));
+    TEST_DO(verify("tensor(x{}){x:(double)}", "double"));
+    TEST_DO(verify("tensor(x{}){x:tensor(x[3])}", "error"));
+    TEST_DO(verify("tensor(x{},y[3]){x:foo,y:2}", "double"));
+    TEST_DO(verify("tensor(x{},y[3]){x:foo}", "tensor(y[3])"));
+    TEST_DO(verify("tensor(x{},y[3]){y:2}", "tensor(x{})"));
+    TEST_DO(verify("tensor<float>(x[3]){x:1}", "double"));
+    TEST_DO(verify("tensor<float>(x[3]){x:double}", "error"));
+    TEST_DO(verify("tensor<float>(x[3]){x:(double)}", "double"));
+    TEST_DO(verify("tensor<float>(x[3]){x:3}", "error"));
+    TEST_DO(verify("tensor<float>(x{}){x:1}", "double"));
+    TEST_DO(verify("tensor<float>(x{}){x:foo}", "double"));
+    TEST_DO(verify("tensor<float>(x{}){x:(double)}", "double"));
+    TEST_DO(verify("tensor<float>(x{}){x:tensor(x[3])}", "error"));
+    TEST_DO(verify("tensor<float>(x{},y[3]){x:foo,y:2}", "double"));
+    TEST_DO(verify("tensor<float>(x{},y[3]){x:foo}", "tensor<float>(y[3])"));
+    TEST_DO(verify("tensor<float>(x{},y[3]){y:2}", "tensor<float>(x{})"));
+}
+
 TEST("require that tensor concat resolves correct type") {
     TEST_DO(verify("concat(double,double,x)", "tensor(x[2])"));
     TEST_DO(verify("concat(tensor(x[2]),tensor(x[3]),x)", "tensor(x[5])"));
diff --git a/eval/src/tests/eval/tensor_function/tensor_function_test.cpp b/eval/src/tests/eval/tensor_function/tensor_function_test.cpp
index c30438b551d..1e59b35a01f 100644
--- a/eval/src/tests/eval/tensor_function/tensor_function_test.cpp
+++ b/eval/src/tests/eval/tensor_function/tensor_function_test.cpp
@@ -6,6 +6,7 @@
 #include <vespa/eval/eval/tensor_function.h>
 #include <vespa/eval/eval/value_type.h>
 #include <vespa/vespalib/util/stash.h>
+#include <vespa/vespalib/util/stringfmt.h>
 #include <map>
 
 using namespace vespalib;
@@ -49,12 +50,27 @@ struct EvalCtx {
     Value::UP make_false() {
         return engine.from_spec(TensorSpec("double").add({}, 0.0));
     }
-    Value::UP make_simple_vector() {
+    Value::UP make_vector(std::initializer_list<double> cells, vespalib::string dim = "x", bool mapped = false) {
+        vespalib::string type_spec = mapped
+                                     ? make_string("tensor(%s{})", dim.c_str())
+                                     : make_string("tensor(%s[%zu])", dim.c_str(), cells.size());
+        TensorSpec spec(type_spec);
+        size_t idx = 0;
+        for (double cell_value: cells) {
+            TensorSpec::Label label = mapped
+                                      ? TensorSpec::Label(make_string("%zu", idx++))
+                                      : TensorSpec::Label(idx++);
+            spec.add({{dim, label}}, cell_value);
+        }
+        return engine.from_spec(spec);
+    }
+    Value::UP make_mixed_tensor(double a, double b, double c, double d) {
         return engine.from_spec(
-                TensorSpec("tensor(x[3])")
-                .add({{"x",0}}, 1)
-                .add({{"x",1}}, 2)
-                .add({{"x",2}}, 3));
+                TensorSpec("tensor(x{},y[2])")
+                .add({{"x", "foo"}, {"y", 0}}, a)
+                .add({{"x", "foo"}, {"y", 1}}, b)
+                .add({{"x", "bar"}, {"y", 0}}, c)
+                .add({{"x", "bar"}, {"y", 1}}, d));
     }
     Value::UP make_tensor_matrix_first_half() {
         return engine.from_spec(
@@ -240,7 +256,7 @@ TEST("require that tensor create works") {
     size_t a_id = ctx.add_tensor(ctx.make_double(1.0));
     size_t b_id = ctx.add_tensor(ctx.make_double(2.0));
     Value::UP my_const = ctx.make_double(3.0);
-    Value::UP expect = ctx.make_simple_vector();
+    Value::UP expect = ctx.make_vector({1.0, 2.0, 3.0});
     const auto &a = inject(ValueType::from_spec("double"), a_id, ctx.stash);
     const auto &b = inject(ValueType::from_spec("double"), b_id, ctx.stash);
     const auto &c = const_value(*my_const, ctx.stash);
@@ -257,6 +273,56 @@ TEST("require that tensor create works") {
     TEST_DO(verify_equal(*expect, ctx.eval(prog)));
 }
 
+TEST("require that single value tensor peek works") {
+    EvalCtx ctx(SimpleTensorEngine::ref());
+    size_t a_id = ctx.add_tensor(ctx.make_double(1.0));
+    Value::UP my_const = ctx.make_mixed_tensor(1.0, 2.0, 3.0, 4.0);
+    Value::UP expect = ctx.make_vector({2.0, 3.0, 0.0});
+    const auto &a = inject(ValueType::from_spec("double"), a_id, ctx.stash);
+    const auto &t = const_value(*my_const, ctx.stash);
+    const auto &peek1 = peek(t, {{"x", "foo"}, {"y", a}}, ctx.stash);
+    const auto &peek2 = peek(t, {{"x", "bar"}, {"y", size_t(0)}}, ctx.stash);
+    const auto &peek3 = peek(t, {{"x", "bar"}, {"y", size_t(1000)}}, ctx.stash);
+    const auto &fun = create(ValueType::from_spec("tensor(x[3])"),
+                             {
+                                 {{{"x", 0}}, peek1},
+                                 {{{"x", 1}}, peek2},
+                                 {{{"x", 2}}, peek3}
+                             },
+                             ctx.stash);
+    EXPECT_TRUE(fun.result_is_mutable());
+    EXPECT_EQUAL(expect->type(), fun.result_type());
+    const auto &prog = ctx.compile(fun);
+    TEST_DO(verify_equal(*expect, ctx.eval(prog)));
+}
+
+TEST("require that tensor subspace tensor peek works") {
+    EvalCtx ctx(SimpleTensorEngine::ref());
+    Value::UP my_const = ctx.make_mixed_tensor(1.0, 2.0, 3.0, 4.0);
+    Value::UP expect = ctx.make_vector({3.0, 4.0}, "y");
+    const auto &t = const_value(*my_const, ctx.stash);
+    const auto &fun = peek(t, {{"x", "bar"}}, ctx.stash);
+    EXPECT_TRUE(fun.result_is_mutable());
+    EXPECT_EQUAL(expect->type(), fun.result_type());
+    const auto &prog = ctx.compile(fun);
+    TEST_DO(verify_equal(*expect, ctx.eval(prog)));
+}
+
+TEST("require that automatic string conversion tensor peek works") {
+    EvalCtx ctx(SimpleTensorEngine::ref());
+    size_t a_id = ctx.add_tensor(ctx.make_double(1.0));
+    Value::UP my_const = ctx.make_vector({1.0, 2.0, 3.0}, "x", true);
+    const auto &a = inject(ValueType::from_spec("double"), a_id, ctx.stash);
+    const auto &t = const_value(*my_const, ctx.stash);
+    const auto &fun = peek(t, {{"x", a}}, ctx.stash);
+    EXPECT_TRUE(fun.result_is_mutable());
+    EXPECT_TRUE(fun.result_type().is_double());
+    const auto &prog = ctx.compile(fun);
+    const Value &result = ctx.eval(prog);
+    EXPECT_TRUE(result.is_double());
+    EXPECT_EQUAL(2.0, result.as_double());
+}
+
 TEST("require that tensor rename works") {
     EvalCtx ctx(SimpleTensorEngine::ref());
     size_t a_id = ctx.add_tensor(ctx.make_tensor_matrix());
@@ -365,20 +431,20 @@ TEST("require that push_children works") {
 
 TEST("require that tensor function can be dumped for debugging") {
     Stash stash;
-    auto my_value_1 = stash.create<DoubleValue>(5.0);
-    auto my_value_2 = stash.create<DoubleValue>(1.0);
+    auto my_value_1 = stash.create<DoubleValue>(1.0);
+    auto my_value_2 = stash.create<DoubleValue>(2.0);
     //-------------------------------------------------------------------------
     const auto &x5 = inject(ValueType::from_spec("tensor(x[5])"), 0, stash);
     const auto &mapped_x5 = map(x5, operation::Relu::f, stash);
     const auto &const_1 = const_value(my_value_1, stash);
     const auto &joined_x5 = join(mapped_x5, const_1, operation::Mul::f, stash);
     //-------------------------------------------------------------------------
-    const auto &x2_0 = const_value(my_value_1, stash);
-    const auto &x2_1 = const_value(my_value_2, stash);
+    const auto &peek1 = peek(x5, {{"x", const_1}}, stash);
+    const auto &peek2 = peek(x5, {{"x", size_t(2)}}, stash);
     const auto &x2 = create(ValueType::from_spec("tensor(x[2])"),
                             {
-                                {{{"x", 0}}, x2_0},
-                                {{{"x", 1}}, x2_1}
+                                {{{"x", 0}}, peek1},
+                                {{{"x", 1}}, peek2}
                             }, stash);
     const auto &a3y10 = inject(ValueType::from_spec("tensor(a[3],y[10])"), 2, stash);
     const auto &a3 = reduce(a3y10, Aggr::SUM, {"y"}, stash);