pow 2/3 as map for dense tensors

- let tensor_function::Node be an implementation detail and not an
  interface requirement.

15 files changed, 279 insertions, 57 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index 67f9fa19dc0..eb09156445d 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -38,6 +38,7 @@ vespa_define_module(
     src/tests/tensor/dense_matmul_function
     src/tests/tensor/dense_multi_matmul_function
     src/tests/tensor/dense_number_join_function
+    src/tests/tensor/dense_pow_as_map_optimizer
     src/tests/tensor/dense_remove_dimension_optimizer
     src/tests/tensor/dense_replace_type_function
     src/tests/tensor/dense_simple_join_function
diff --git a/eval/src/tests/eval/inline_operation/inline_operation_test.cpp b/eval/src/tests/eval/inline_operation/inline_operation_test.cpp
index 8895bd4bcbd..fe9396398da 100644
--- a/eval/src/tests/eval/inline_operation/inline_operation_test.cpp
+++ b/eval/src/tests/eval/inline_operation/inline_operation_test.cpp
@@ -68,6 +68,18 @@ TEST(InlineOperationTest, op1_lambdas_are_recognized) {
     //-------------------------------------------
     EXPECT_EQ(as_op1("1/a"),        &Inv::f);
     EXPECT_EQ(as_op1("1.0/a"),      &Inv::f);
+    EXPECT_EQ(as_op1("a*a"),        &Square::f);
+    EXPECT_EQ(as_op1("a^2"),        &Square::f);
+    EXPECT_EQ(as_op1("a^2.0"),      &Square::f);
+    EXPECT_EQ(as_op1("pow(a,2)"),   &Square::f);
+    EXPECT_EQ(as_op1("pow(a,2.0)"), &Square::f);
+    EXPECT_EQ(as_op1("a*a*a"),      &Cube::f);
+    EXPECT_EQ(as_op1("(a*a)*a"),    &Cube::f);
+    EXPECT_EQ(as_op1("a*(a*a)"),    &Cube::f);
+    EXPECT_EQ(as_op1("a^3"),        &Cube::f);
+    EXPECT_EQ(as_op1("a^3.0"),      &Cube::f);
+    EXPECT_EQ(as_op1("pow(a,3)"),   &Cube::f);
+    EXPECT_EQ(as_op1("pow(a,3.0)"), &Cube::f);
 }
 
 TEST(InlineOperationTest, op1_lambdas_are_recognized_with_different_parameter_names) {
@@ -76,7 +88,7 @@ TEST(InlineOperationTest, op1_lambdas_are_recognized_with_different_parameter_na
 }
 
 TEST(InlineOperationTest, non_op1_lambdas_are_not_recognized) {
-    EXPECT_FALSE(lookup_op1(*Function::parse({"a"}, "a*a")).has_value());
+    EXPECT_FALSE(lookup_op1(*Function::parse({"a"}, "a*a+3")).has_value());
     EXPECT_FALSE(lookup_op1(*Function::parse({"a", "b"}, "a+b")).has_value());
 }
 
@@ -170,6 +182,19 @@ TEST(InlineOperationTest, parameter_swap_wrapper_works) {
 
 //-----------------------------------------------------------------------------
 
+TEST(InlineOperationTest, op1_cube_is_inlined) {
+    TypifyOp1::resolve(Cube::f, [](auto t)
+                       {
+                           using T = typename decltype(t)::type;
+                           bool type_ok = std::is_same_v<T,InlineOp1<Cube>>;
+                           op1_t ref = Cube::f;
+                           EXPECT_TRUE(type_ok);
+                           test_op1<T>(ref, 2.0, 8.0);
+                           test_op1<T>(ref, 3.0, 27.0);
+                           test_op1<T>(ref, 7.0, 343.0);
+                       });
+}
+
 TEST(InlineOperationTest, op1_exp_is_inlined) {
     TypifyOp1::resolve(Exp::f, [](auto t)
                        {
@@ -209,6 +234,19 @@ TEST(InlineOperationTest, op1_sqrt_is_inlined) {
                        });
 }
 
+TEST(InlineOperationTest, op1_square_is_inlined) {
+    TypifyOp1::resolve(Square::f, [](auto t)
+                       {
+                           using T = typename decltype(t)::type;
+                           bool type_ok = std::is_same_v<T,InlineOp1<Square>>;
+                           op1_t ref = Square::f;
+                           EXPECT_TRUE(type_ok);
+                           test_op1<T>(ref, 2.0, 4.0);
+                           test_op1<T>(ref, 3.0, 9.0);
+                           test_op1<T>(ref, 7.0, 49.0);
+                       });
+}
+
 TEST(InlineOperationTest, op1_tanh_is_inlined) {
     TypifyOp1::resolve(Tanh::f, [](auto t)
                        {
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 1eb9912abd2..b205205f52e 100644
--- a/eval/src/tests/eval/tensor_function/tensor_function_test.cpp
+++ b/eval/src/tests/eval/tensor_function/tensor_function_test.cpp
@@ -38,7 +38,7 @@ struct EvalCtx {
         ictx = std::make_unique<InterpretedFunction::Context>(*ifun);
         return ifun->eval(*ictx, SimpleObjectParams(params));
     }
-    const TensorFunction &compile(const tensor_function::Node &expr) {
+    const TensorFunction &compile(const TensorFunction &expr) {
         return engine.optimize(expr, stash);
     }
     Value::UP make_double(double value) {
@@ -391,15 +391,15 @@ TEST("require that if_node works") {
 
 TEST("require that if_node result is mutable only when both children produce mutable results") {
     Stash stash;
-    const Node &cond = inject(DoubleValue::double_type(), 0, stash);
-    const Node &a = inject(ValueType::from_spec("tensor(x[2])"), 0, stash);
-    const Node &b = inject(ValueType::from_spec("tensor(x[3])"), 0, stash);
-    const Node &c = inject(ValueType::from_spec("tensor(x[5])"), 0, stash);
-    const Node &tmp = concat(a, b, "x", stash); // will be mutable
-    const Node &if_con_con = if_node(cond, c, c, stash);
-    const Node &if_mut_con = if_node(cond, tmp, c, stash);
-    const Node &if_con_mut = if_node(cond, c, tmp, stash);
-    const Node &if_mut_mut = if_node(cond, tmp, tmp, stash);
+    const TensorFunction &cond = inject(DoubleValue::double_type(), 0, stash);
+    const TensorFunction &a = inject(ValueType::from_spec("tensor(x[2])"), 0, stash);
+    const TensorFunction &b = inject(ValueType::from_spec("tensor(x[3])"), 0, stash);
+    const TensorFunction &c = inject(ValueType::from_spec("tensor(x[5])"), 0, stash);
+    const TensorFunction &tmp = concat(a, b, "x", stash); // will be mutable
+    const TensorFunction &if_con_con = if_node(cond, c, c, stash);
+    const TensorFunction &if_mut_con = if_node(cond, tmp, c, stash);
+    const TensorFunction &if_con_mut = if_node(cond, c, tmp, stash);
+    const TensorFunction &if_mut_mut = if_node(cond, tmp, tmp, stash);
     EXPECT_EQUAL(if_con_con.result_type(), c.result_type());
     EXPECT_EQUAL(if_con_mut.result_type(), c.result_type());
     EXPECT_EQUAL(if_mut_con.result_type(), c.result_type());
@@ -412,13 +412,13 @@ TEST("require that if_node result is mutable only when both children produce mut
 
 TEST("require that if_node gets expected result type") {
     Stash stash;
-    const Node &a = inject(DoubleValue::double_type(), 0, stash);
-    const Node &b = inject(ValueType::from_spec("tensor(x[2])"), 0, stash);
-    const Node &c = inject(ValueType::from_spec("tensor(x[3])"), 0, stash);
-    const Node &d = inject(ValueType::from_spec("error"), 0, stash);
-    const Node &if_same = if_node(a, b, b, stash);
-    const Node &if_different = if_node(a, b, c, stash);
-    const Node &if_with_error = if_node(a, b, d, stash);
+    const TensorFunction &a = inject(DoubleValue::double_type(), 0, stash);
+    const TensorFunction &b = inject(ValueType::from_spec("tensor(x[2])"), 0, stash);
+    const TensorFunction &c = inject(ValueType::from_spec("tensor(x[3])"), 0, stash);
+    const TensorFunction &d = inject(ValueType::from_spec("error"), 0, stash);
+    const TensorFunction &if_same = if_node(a, b, b, stash);
+    const TensorFunction &if_different = if_node(a, b, c, stash);
+    const TensorFunction &if_with_error = if_node(a, b, d, stash);
     EXPECT_EQUAL(if_same.result_type(), ValueType::from_spec("tensor(x[2])"));
     EXPECT_EQUAL(if_different.result_type(), ValueType::from_spec("error"));
     EXPECT_EQUAL(if_with_error.result_type(), ValueType::from_spec("error"));
@@ -426,10 +426,10 @@ TEST("require that if_node gets expected result type") {
 
 TEST("require that push_children works") {
     Stash stash;
-    std::vector<Node::Child::CREF> refs;
-    const Node &a = inject(DoubleValue::double_type(), 0, stash);
-    const Node &b = inject(DoubleValue::double_type(), 1, stash);
-    const Node &c = const_value(stash.create<DoubleValue>(1.0), stash);
+    std::vector<TensorFunction::Child::CREF> refs;
+    const TensorFunction &a = inject(DoubleValue::double_type(), 0, stash);
+    const TensorFunction &b = inject(DoubleValue::double_type(), 1, stash);
+    const TensorFunction &c = const_value(stash.create<DoubleValue>(1.0), stash);
     a.push_children(refs);
     b.push_children(refs);
     c.push_children(refs);
diff --git a/eval/src/tests/tensor/dense_pow_as_map_optimizer/CMakeLists.txt b/eval/src/tests/tensor/dense_pow_as_map_optimizer/CMakeLists.txt
new file mode 100644
index 00000000000..cf5103f8b4f
--- /dev/null
+++ b/eval/src/tests/tensor/dense_pow_as_map_optimizer/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(eval_dense_pow_as_map_optimizer_test_app TEST
+    SOURCES
+    dense_pow_as_map_optimizer_test.cpp
+    DEPENDS
+    vespaeval
+    gtest
+)
+vespa_add_test(NAME eval_dense_pow_as_map_optimizer_test_app COMMAND eval_dense_pow_as_map_optimizer_test_app)
diff --git a/eval/src/tests/tensor/dense_pow_as_map_optimizer/dense_pow_as_map_optimizer_test.cpp b/eval/src/tests/tensor/dense_pow_as_map_optimizer/dense_pow_as_map_optimizer_test.cpp
new file mode 100644
index 00000000000..38d9ac8aeef
--- /dev/null
+++ b/eval/src/tests/tensor/dense_pow_as_map_optimizer/dense_pow_as_map_optimizer_test.cpp
@@ -0,0 +1,92 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/eval/eval/tensor_function.h>
+#include <vespa/eval/tensor/default_tensor_engine.h>
+#include <vespa/eval/tensor/dense/dense_simple_map_function.h>
+#include <vespa/eval/eval/test/eval_fixture.h>
+#include <vespa/eval/eval/test/tensor_model.hpp>
+#include <vespa/vespalib/gtest/gtest.h>
+
+using namespace vespalib::eval::operation;
+using namespace vespalib::eval::tensor_function;
+using namespace vespalib::eval::test;
+using namespace vespalib::eval;
+using namespace vespalib::tensor;
+//using namespace vespalib;
+
+const TensorEngine &prod_engine = DefaultTensorEngine::ref();
+
+EvalFixture::ParamRepo make_params() {
+    return EvalFixture::ParamRepo()
+        .add("a", spec(1.5))
+        .add("b", spec(2.5))
+        .add("sparse", spec({x({"a"})}, N()))
+        .add("mixed", spec({x({"a"}),y(5)}, N()))
+        .add_matrix("x", 5, "y", 3);
+}
+EvalFixture::ParamRepo param_repo = make_params();
+
+void verify_optimized(const vespalib::string &expr, op1_t op1, bool inplace = false) {
+    EvalFixture slow_fixture(prod_engine, expr, param_repo, false);
+    EvalFixture fixture(prod_engine, expr, param_repo, true, true);
+    EXPECT_EQ(fixture.result(), EvalFixture::ref(expr, param_repo));
+    EXPECT_EQ(fixture.result(), slow_fixture.result());
+    auto info = fixture.find_all<DenseSimpleMapFunction>();
+    ASSERT_EQ(info.size(), 1u);
+    EXPECT_TRUE(info[0]->result_is_mutable());
+    EXPECT_EQ(info[0]->function(), op1);
+    EXPECT_EQ(info[0]->inplace(), inplace);
+    ASSERT_EQ(fixture.num_params(), 1);
+    if (inplace) {
+        EXPECT_EQ(fixture.get_param(0), fixture.result());
+    } else {
+        EXPECT_TRUE(!(fixture.get_param(0) == fixture.result()));
+    }
+}
+
+void verify_not_optimized(const vespalib::string &expr) {
+    EvalFixture slow_fixture(prod_engine, expr, param_repo, false);
+    EvalFixture fixture(prod_engine, expr, param_repo, true);
+    EXPECT_EQ(fixture.result(), EvalFixture::ref(expr, param_repo));
+    EXPECT_EQ(fixture.result(), slow_fixture.result());
+    auto info = fixture.find_all<Map>();
+    EXPECT_TRUE(info.empty());
+}
+
+TEST(PowAsMapTest, squared_dense_tensor_is_optimized) {
+    verify_optimized("x5y3^2.0", Square::f);
+    verify_optimized("pow(x5y3,2.0)", Square::f);
+    verify_optimized("join(x5y3,2.0,f(x,y)(x^y))", Square::f);
+    verify_optimized("join(x5y3,2.0,f(x,y)(pow(x,y)))", Square::f);
+    verify_optimized("join(x5y3f,2.0,f(x,y)(pow(x,y)))", Square::f);
+    verify_optimized("join(@x5y3,2.0,f(x,y)(pow(x,y)))", Square::f, true);
+    verify_optimized("join(@x5y3f,2.0,f(x,y)(pow(x,y)))", Square::f, true);
+}
+
+TEST(PowAsMapTest, cubed_dense_tensor_is_optimized) {
+    verify_optimized("x5y3^3.0", Cube::f);
+    verify_optimized("pow(x5y3,3.0)", Cube::f);
+    verify_optimized("join(x5y3,3.0,f(x,y)(x^y))", Cube::f);
+    verify_optimized("join(x5y3,3.0,f(x,y)(pow(x,y)))", Cube::f);
+    verify_optimized("join(x5y3f,3.0,f(x,y)(pow(x,y)))", Cube::f);
+    verify_optimized("join(@x5y3,3.0,f(x,y)(pow(x,y)))", Cube::f, true);
+    verify_optimized("join(@x5y3f,3.0,f(x,y)(pow(x,y)))", Cube::f, true);
+}
+
+TEST(PowAsMapTest, hypercubed_dense_tensor_is_not_optimized) {
+    verify_not_optimized("join(x5y3,4.0,f(x,y)(pow(x,y)))");
+}
+
+TEST(PowAsMapTest, scalar_join_is_not_optimized) {
+    verify_not_optimized("join(a,2.0,f(x,y)(pow(x,y)))");
+}
+
+TEST(PowAsMapTest, sparse_join_is_not_optimized) {
+    verify_not_optimized("join(sparse,2.0,f(x,y)(pow(x,y)))");
+}
+
+TEST(PowAsMapTest, mixed_join_is_not_optimized) {
+    verify_not_optimized("join(mixed,2.0,f(x,y)(pow(x,y)))");
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/eval/src/vespa/eval/eval/inline_operation.h b/eval/src/vespa/eval/eval/inline_operation.h
index fccf1874242..21516c4d94e 100644
--- a/eval/src/vespa/eval/eval/inline_operation.h
+++ b/eval/src/vespa/eval/eval/inline_operation.h
@@ -17,6 +17,10 @@ struct CallOp1 {
 };
 
 template <typename T> struct InlineOp1;
+template <> struct InlineOp1<Cube> {
+    InlineOp1(op1_t) {}
+    template <typename A> constexpr auto operator()(A a) const { return (a * a * a); }
+};
 template <> struct InlineOp1<Exp> {
     InlineOp1(op1_t) {}
     template <typename A> constexpr auto operator()(A a) const { return exp(a); }
@@ -29,6 +33,10 @@ template <> struct InlineOp1<Sqrt> {
     InlineOp1(op1_t) {}
     template <typename A> constexpr auto operator()(A a) const { return std::sqrt(a); }
 };
+template <> struct InlineOp1<Square> {
+    InlineOp1(op1_t) {}
+    template <typename A> constexpr auto operator()(A a) const { return (a * a); }
+};
 template <> struct InlineOp1<Tanh> {
     InlineOp1(op1_t) {}
     template <typename A> constexpr auto operator()(A a) const { return std::tanh(a); }
@@ -37,12 +45,16 @@ template <> struct InlineOp1<Tanh> {
 struct TypifyOp1 {
     template <typename T> using Result = TypifyResultType<T>;
     template <typename F> static decltype(auto) resolve(op1_t value, F &&f) {
-        if (value == Exp::f) {
+        if (value == Cube::f) {
+            return f(Result<InlineOp1<Cube>>());
+        } else if (value == Exp::f) {
             return f(Result<InlineOp1<Exp>>());
         } else if (value == Inv::f) {
             return f(Result<InlineOp1<Inv>>());
         } else if (value == Sqrt::f) {
             return f(Result<InlineOp1<Sqrt>>());
+        } else if (value == Square::f) {
+            return f(Result<InlineOp1<Square>>());
         } else if (value == Tanh::f) {
             return f(Result<InlineOp1<Tanh>>());
         } else {
diff --git a/eval/src/vespa/eval/eval/make_tensor_function.cpp b/eval/src/vespa/eval/eval/make_tensor_function.cpp
index 02d16caae6b..e80633b5c41 100644
--- a/eval/src/vespa/eval/eval/make_tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/make_tensor_function.cpp
@@ -22,7 +22,7 @@ struct TensorFunctionBuilder : public NodeVisitor, public NodeTraverser {
     Stash              &stash;
     const TensorEngine &tensor_engine;
     const NodeTypes    &types;
-    std::vector<tensor_function::Node::CREF> stack;
+    std::vector<TensorFunction::CREF> stack;
 
     TensorFunctionBuilder(Stash &stash_in, const TensorEngine &tensor_engine_in, const NodeTypes &types_in)
         : stash(stash_in), tensor_engine(tensor_engine_in), types(types_in), stack() {}
@@ -94,7 +94,7 @@ struct TensorFunctionBuilder : public NodeVisitor, public NodeTraverser {
 
     void make_create(const TensorCreate &node) {
         assert(stack.size() >= node.num_children());
-        std::map<TensorSpec::Address, tensor_function::Node::CREF> spec;
+        std::map<TensorSpec::Address, TensorFunction::CREF> spec;
         for (size_t idx = node.num_children(); idx-- > 0; ) {
             spec.emplace(node.get_child_address(idx), stack.back());
             stack.pop_back();
@@ -115,8 +115,8 @@ struct TensorFunctionBuilder : public NodeVisitor, public NodeTraverser {
 
     void make_peek(const TensorPeek &node) {
         assert(stack.size() >= node.num_children());
-        const tensor_function::Node &param = stack[stack.size()-node.num_children()];
-        std::map<vespalib::string, std::variant<TensorSpec::Label, tensor_function::Node::CREF>> spec;
+        const TensorFunction &param = stack[stack.size()-node.num_children()];
+        std::map<vespalib::string, std::variant<TensorSpec::Label, TensorFunction::CREF>> spec;
         for (auto pos = node.dim_list().rbegin(); pos != node.dim_list().rend(); ++pos) {
             if (pos->second.is_expr()) {
                 spec.emplace(pos->first, stack.back());
diff --git a/eval/src/vespa/eval/eval/operation.cpp b/eval/src/vespa/eval/eval/operation.cpp
index bbd37ab68b2..fa8be4d20bc 100644
--- a/eval/src/vespa/eval/eval/operation.cpp
+++ b/eval/src/vespa/eval/eval/operation.cpp
@@ -50,7 +50,9 @@ double Relu::f(double a) { return std::max(a, 0.0); }
 double Sigmoid::f(double a) { return 1.0 / (1.0 + std::exp(-1.0 * a)); }
 double Elu::f(double a) { return (a < 0) ? std::exp(a) - 1 : a; }
 //-----------------------------------------------------------------------------
-double Inv::f(double a) { return (1 / a); }
+double Inv::f(double a) { return (1.0 / a); }
+double Square::f(double a) { return (a * a); }
+double Cube::f(double a) { return (a * a * a); }
 
 namespace {
 
@@ -106,6 +108,13 @@ std::map<vespalib::string,op1_t> make_op1_map() {
     add_op1(map, "elu(a)",     Elu::f);
     //-------------------------------------
     add_op1(map, "1/a",        Inv::f);
+    add_op1(map, "a*a",        Square::f);
+    add_op1(map, "a^2",        Square::f);
+    add_op1(map, "pow(a,2)",   Square::f);
+    add_op1(map, "(a*a)*a",    Cube::f);
+    add_op1(map, "a*(a*a)",    Cube::f);
+    add_op1(map, "a^3",        Cube::f);
+    add_op1(map, "pow(a,3)",   Cube::f);
     return map;
 }
 
diff --git a/eval/src/vespa/eval/eval/operation.h b/eval/src/vespa/eval/eval/operation.h
index b00bb5e26fc..02d3322f867 100644
--- a/eval/src/vespa/eval/eval/operation.h
+++ b/eval/src/vespa/eval/eval/operation.h
@@ -50,6 +50,8 @@ struct Sigmoid { static double f(double a); };
 struct Elu { static double f(double a); };
 //-----------------------------------------------------------------------------
 struct Inv { static double f(double a); };
+struct Square { static double f(double a); };
+struct Cube { static double f(double a); };
 
 using op1_t = double (*)(double);
 using op2_t = double (*)(double, double);
diff --git a/eval/src/vespa/eval/eval/tensor_function.cpp b/eval/src/vespa/eval/eval/tensor_function.cpp
index 1aa18417b87..85079b7a8e3 100644
--- a/eval/src/vespa/eval/eval/tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/tensor_function.cpp
@@ -545,48 +545,48 @@ If::visit_children(vespalib::ObjectVisitor &visitor) const
 
 //-----------------------------------------------------------------------------
 
-const Node &const_value(const Value &value, Stash &stash) {
+const TensorFunction &const_value(const Value &value, Stash &stash) {
     return stash.create<ConstValue>(value);
 }
 
-const Node &inject(const ValueType &type, size_t param_idx, Stash &stash) {
+const TensorFunction &inject(const ValueType &type, size_t param_idx, Stash &stash) {
     return stash.create<Inject>(type, param_idx);
 }
 
-const Node &reduce(const Node &child, Aggr aggr, const std::vector<vespalib::string> &dimensions, Stash &stash) {
+const TensorFunction &reduce(const TensorFunction &child, Aggr aggr, const std::vector<vespalib::string> &dimensions, Stash &stash) {
     ValueType result_type = child.result_type().reduce(dimensions);
     return stash.create<Reduce>(result_type, child, aggr, dimensions);
 }
 
-const Node &map(const Node &child, map_fun_t function, Stash &stash) {
+const TensorFunction &map(const TensorFunction &child, map_fun_t function, Stash &stash) {
     ValueType result_type = child.result_type();
     return stash.create<Map>(result_type, child, function);
 }
 
-const Node &join(const Node &lhs, const Node &rhs, join_fun_t function, Stash &stash) {
+const TensorFunction &join(const TensorFunction &lhs, const TensorFunction &rhs, join_fun_t function, Stash &stash) {
     ValueType result_type = ValueType::join(lhs.result_type(), rhs.result_type());
     return stash.create<Join>(result_type, lhs, rhs, function);
 }
 
-const Node &merge(const Node &lhs, const Node &rhs, join_fun_t function, Stash &stash) {
+const TensorFunction &merge(const TensorFunction &lhs, const TensorFunction &rhs, join_fun_t function, Stash &stash) {
     ValueType result_type = ValueType::merge(lhs.result_type(), rhs.result_type());
     return stash.create<Merge>(result_type, lhs, rhs, function);
 }
 
-const Node &concat(const Node &lhs, const Node &rhs, const vespalib::string &dimension, Stash &stash) {
+const TensorFunction &concat(const TensorFunction &lhs, const TensorFunction &rhs, const vespalib::string &dimension, Stash &stash) {
     ValueType result_type = ValueType::concat(lhs.result_type(), rhs.result_type(), dimension);
     return stash.create<Concat>(result_type, lhs, rhs, dimension);
 }
 
-const Node &create(const ValueType &type, const std::map<TensorSpec::Address,Node::CREF> &spec, Stash &stash) {
+const TensorFunction &create(const ValueType &type, const std::map<TensorSpec::Address,TensorFunction::CREF> &spec, Stash &stash) {
     return stash.create<Create>(type, spec);
 }
 
-const Node &lambda(const ValueType &type, const std::vector<size_t> &bindings, const Function &function, NodeTypes node_types, Stash &stash) {
+const TensorFunction &lambda(const ValueType &type, const std::vector<size_t> &bindings, const Function &function, NodeTypes node_types, Stash &stash) {
     return stash.create<Lambda>(type, bindings, function, std::move(node_types));
 }
 
-const Node &peek(const Node &param, const std::map<vespalib::string, std::variant<TensorSpec::Label, Node::CREF>> &spec, Stash &stash) {
+const TensorFunction &peek(const TensorFunction &param, const std::map<vespalib::string, std::variant<TensorSpec::Label, TensorFunction::CREF>> &spec, Stash &stash) {
     std::vector<vespalib::string> dimensions;
     for (const auto &dim_spec: spec) {
         dimensions.push_back(dim_spec.first);
@@ -595,12 +595,12 @@ const Node &peek(const Node &param, const std::map<vespalib::string, std::varian
     return stash.create<Peek>(result_type, param, spec);
 }
 
-const Node &rename(const Node &child, const std::vector<vespalib::string> &from, const std::vector<vespalib::string> &to, Stash &stash) {
+const TensorFunction &rename(const TensorFunction &child, const std::vector<vespalib::string> &from, const std::vector<vespalib::string> &to, Stash &stash) {
     ValueType result_type = child.result_type().rename(from, to);
     return stash.create<Rename>(result_type, child, from, to);
 }
 
-const Node &if_node(const Node &cond, const Node &true_child, const Node &false_child, Stash &stash) {
+const TensorFunction &if_node(const TensorFunction &cond, const TensorFunction &true_child, const TensorFunction &false_child, Stash &stash) {
     ValueType result_type = ValueType::either(true_child.result_type(), false_child.result_type());
     return stash.create<If>(result_type, cond, true_child, false_child);
 }
diff --git a/eval/src/vespa/eval/eval/tensor_function.h b/eval/src/vespa/eval/eval/tensor_function.h
index 6743f37eeb1..20631108775 100644
--- a/eval/src/vespa/eval/eval/tensor_function.h
+++ b/eval/src/vespa/eval/eval/tensor_function.h
@@ -52,6 +52,7 @@ class Tensor;
  **/
 struct TensorFunction
 {
+    using CREF = std::reference_wrapper<const TensorFunction>;
     TensorFunction(const TensorFunction &) = delete;
     TensorFunction &operator=(const TensorFunction &) = delete;
     TensorFunction(TensorFunction &&) = delete;
@@ -132,7 +133,6 @@ class Node : public TensorFunction
 private:
     ValueType _result_type;
 public:
-    using CREF = std::reference_wrapper<const Node>;
     Node(const ValueType &result_type_in) : _result_type(result_type_in) {}
     const ValueType &result_type() const final override { return _result_type; }
 };
@@ -310,7 +310,7 @@ class Create : public Node
 private:
     std::map<TensorSpec::Address, Child> _spec;
 public:
-    Create(const ValueType &result_type_in, const std::map<TensorSpec::Address, Node::CREF> &spec_in)
+    Create(const ValueType &result_type_in, const std::map<TensorSpec::Address, TensorFunction::CREF> &spec_in)
         : Super(result_type_in), _spec()
     {
         for (const auto &cell: spec_in) {
@@ -359,8 +359,8 @@ private:
     Child _param;
     std::map<vespalib::string, MyLabel> _spec;
 public:
-    Peek(const ValueType &result_type_in, const Node &param,
-         const std::map<vespalib::string, std::variant<TensorSpec::Label, Node::CREF>> &spec)
+    Peek(const ValueType &result_type_in, const TensorFunction &param,
+         const std::map<vespalib::string, std::variant<TensorSpec::Label, TensorFunction::CREF>> &spec)
         : Super(result_type_in), _param(param), _spec()
     {
         for (const auto &dim: spec) {
@@ -369,7 +369,7 @@ public:
                            [&](const TensorSpec::Label &label) {
                                _spec.emplace(dim.first, label);
                            },
-                           [&](const Node::CREF &ref) {
+                           [&](const TensorFunction::CREF &ref) {
                                _spec.emplace(dim.first, ref.get());
                            }
                        }, dim.second);
@@ -432,18 +432,18 @@ public:
 
 //-----------------------------------------------------------------------------
 
-const Node &const_value(const Value &value, Stash &stash);
-const Node &inject(const ValueType &type, size_t param_idx, Stash &stash);
-const Node &reduce(const Node &child, Aggr aggr, const std::vector<vespalib::string> &dimensions, Stash &stash);
-const Node &map(const Node &child, map_fun_t function, Stash &stash);
-const Node &join(const Node &lhs, const Node &rhs, join_fun_t function, Stash &stash);
-const Node &merge(const Node &lhs, const Node &rhs, join_fun_t function, Stash &stash);
-const Node &concat(const Node &lhs, const Node &rhs, const vespalib::string &dimension, Stash &stash);
-const Node &create(const ValueType &type, const std::map<TensorSpec::Address, Node::CREF> &spec, Stash &stash);
-const Node &lambda(const ValueType &type, const std::vector<size_t> &bindings, const Function &function, NodeTypes node_types, Stash &stash);
-const Node &peek(const Node &param, const std::map<vespalib::string, std::variant<TensorSpec::Label, Node::CREF>> &spec, Stash &stash);
-const Node &rename(const Node &child, const std::vector<vespalib::string> &from, const std::vector<vespalib::string> &to, Stash &stash);
-const Node &if_node(const Node &cond, const Node &true_child, const Node &false_child, Stash &stash);
+const TensorFunction &const_value(const Value &value, Stash &stash);
+const TensorFunction &inject(const ValueType &type, size_t param_idx, Stash &stash);
+const TensorFunction &reduce(const TensorFunction &child, Aggr aggr, const std::vector<vespalib::string> &dimensions, Stash &stash);
+const TensorFunction &map(const TensorFunction &child, map_fun_t function, Stash &stash);
+const TensorFunction &join(const TensorFunction &lhs, const TensorFunction &rhs, join_fun_t function, Stash &stash);
+const TensorFunction &merge(const TensorFunction &lhs, const TensorFunction &rhs, join_fun_t function, Stash &stash);
+const TensorFunction &concat(const TensorFunction &lhs, const TensorFunction &rhs, const vespalib::string &dimension, Stash &stash);
+const TensorFunction &create(const ValueType &type, const std::map<TensorSpec::Address, TensorFunction::CREF> &spec, Stash &stash);
+const TensorFunction &lambda(const ValueType &type, const std::vector<size_t> &bindings, const Function &function, NodeTypes node_types, Stash &stash);
+const TensorFunction &peek(const TensorFunction &param, const std::map<vespalib::string, std::variant<TensorSpec::Label, TensorFunction::CREF>> &spec, Stash &stash);
+const TensorFunction &rename(const TensorFunction &child, const std::vector<vespalib::string> &from, const std::vector<vespalib::string> &to, Stash &stash);
+const TensorFunction &if_node(const TensorFunction &cond, const TensorFunction &true_child, const TensorFunction &false_child, Stash &stash);
 
 } // namespace vespalib::eval::tensor_function
 } // namespace vespalib::eval
diff --git a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
index 22073143ac0..cb933c74754 100644
--- a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
+++ b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
@@ -20,6 +20,7 @@
 #include "dense/dense_lambda_function.h"
 #include "dense/dense_simple_join_function.h"
 #include "dense/dense_number_join_function.h"
+#include "dense/dense_pow_as_map_optimizer.h"
 #include "dense/dense_simple_map_function.h"
 #include "dense/vector_from_doubles_function.h"
 #include "dense/dense_tensor_create_function.h"
@@ -295,6 +296,7 @@ DefaultTensorEngine::optimize(const TensorFunction &expr, Stash &stash) const
             child.set(DenseLambdaPeekOptimizer::optimize(child.get(), stash));
             child.set(DenseLambdaFunction::optimize(child.get(), stash));
             child.set(DenseFastRenameOptimizer::optimize(child.get(), stash));
+            child.set(DensePowAsMapOptimizer::optimize(child.get(), stash));
             child.set(DenseSimpleMapFunction::optimize(child.get(), stash));
             child.set(DenseSimpleJoinFunction::optimize(child.get(), stash));
             child.set(DenseNumberJoinFunction::optimize(child.get(), stash));
diff --git a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
index 244e288b90a..c33d6f5cbcf 100644
--- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
+++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
@@ -12,6 +12,7 @@ vespa_add_library(eval_tensor_dense OBJECT
     dense_matmul_function.cpp
     dense_multi_matmul_function.cpp
     dense_number_join_function.cpp
+    dense_pow_as_map_optimizer.cpp
     dense_remove_dimension_optimizer.cpp
     dense_replace_type_function.cpp
     dense_simple_join_function.cpp
diff --git a/eval/src/vespa/eval/tensor/dense/dense_pow_as_map_optimizer.cpp b/eval/src/vespa/eval/tensor/dense/dense_pow_as_map_optimizer.cpp
new file mode 100644
index 00000000000..f78c23c80ac
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_pow_as_map_optimizer.cpp
@@ -0,0 +1,38 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "dense_pow_as_map_optimizer.h"
+#include "dense_simple_map_function.h"
+#include <vespa/eval/eval/operation.h>
+
+namespace vespalib::tensor {
+
+using eval::TensorFunction;
+using eval::as;
+
+using namespace eval::tensor_function;
+using namespace eval::operation;
+
+const TensorFunction &
+DensePowAsMapOptimizer::optimize(const TensorFunction &expr, Stash &stash)
+{
+    if (auto join = as<Join>(expr)) {
+        const TensorFunction &lhs = join->lhs();
+        const TensorFunction &rhs = join->rhs();
+        if ((join->function() == Pow::f) &&
+            lhs.result_type().is_dense() &&
+            rhs.result_type().is_double())
+        {
+            if (auto const_value = as<ConstValue>(rhs)) {
+                if (const_value->value().as_double() == 2.0) {
+                    return map(lhs, Square::f, stash);
+                }
+                if (const_value->value().as_double() == 3.0) {
+                    return map(lhs, Cube::f, stash);
+                }
+            }
+        }
+    }
+    return expr;
+}
+
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_pow_as_map_optimizer.h b/eval/src/vespa/eval/tensor/dense/dense_pow_as_map_optimizer.h
new file mode 100644
index 00000000000..4849a10c070
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_pow_as_map_optimizer.h
@@ -0,0 +1,18 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/eval/eval/tensor_function.h>
+
+namespace vespalib::tensor {
+
+/**
+ * Tensor function optimizer for converting join expressions on the
+ * form 'join(tensor,<small integer constant>,f(x,y)(pow(x,y))' to
+ * expressions on the form 'map(tensor,f(x)(x*x...))'.
+ **/
+struct DensePowAsMapOptimizer {
+    static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
+};
+
+} // namespace vespalib::tensor