use new api in tensor function ir

11 files changed, 110 insertions, 110 deletions

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 de97c785916..5b2d0848f64 100644
--- a/eval/src/tests/eval/tensor_function/tensor_function_test.cpp
+++ b/eval/src/tests/eval/tensor_function/tensor_function_test.cpp
@@ -15,11 +15,10 @@ using namespace vespalib::eval::tensor_function;
 struct EvalCtx : TensorFunction::Input {
     const TensorEngine &engine;
     Stash stash;
-    operation::Neg neg;
     ErrorValue error;
     std::map<size_t, Value::UP> tensors;
     EvalCtx(const TensorEngine &engine_in)
-        : engine(engine_in), stash(), neg(), error(), tensors() {}
+        : engine(engine_in), stash(), error(), tensors() {}
     ~EvalCtx() { }
     void add_tensor(std::unique_ptr<Tensor> tensor, size_t id) {
         tensors.emplace(id, std::make_unique<TensorValue>(std::move(tensor)));
@@ -30,10 +29,6 @@ struct EvalCtx : TensorFunction::Input {
         }
         return *tensors.find(id)->second;
     }
-    const UnaryOperation &get_map_operation(size_t id) const override {
-        ASSERT_EQUAL(42u, id);
-        return neg;
-    }
     const Value &eval(const TensorFunction &fun) { return fun.eval(*this, stash); }
     const ValueType type(const Tensor &tensor) const { return engine.type_of(tensor); }
     TensorFunction::UP compile(tensor_function::Node_UP expr) const {
@@ -124,7 +119,7 @@ TEST("require that partial tensor reduction works") {
     EvalCtx ctx(SimpleTensorEngine::ref());
     ctx.add_tensor(ctx.make_tensor_reduce_input(), 1);
     auto expect = ctx.make_tensor_reduce_y_output();
-    auto fun = reduce(inject(ValueType::from_spec("tensor(x[3],y[2])"), 1), operation::Add(), {"y"});
+    auto fun = reduce(inject(ValueType::from_spec("tensor(x[3],y[2])"), 1), Aggr::SUM, {"y"});
     EXPECT_EQUAL(ctx.type(*expect), fun->result_type);
     auto prog = ctx.compile(std::move(fun));
     TEST_DO(verify_equal(*expect, ctx.eval(*prog)));
@@ -133,7 +128,7 @@ TEST("require that partial tensor reduction works") {
 TEST("require that full tensor reduction works") {
     EvalCtx ctx(SimpleTensorEngine::ref());
     ctx.add_tensor(ctx.make_tensor_reduce_input(), 1);
-    auto fun = reduce(inject(ValueType::from_spec("tensor(x[3],y[2])"), 1), operation::Add(), {});
+    auto fun = reduce(inject(ValueType::from_spec("tensor(x[3],y[2])"), 1), Aggr::SUM, {});
     EXPECT_EQUAL(ValueType::from_spec("double"), fun->result_type);
     auto prog = ctx.compile(std::move(fun));
     EXPECT_EQUAL(21.0, ctx.eval(*prog).as_double());
@@ -143,20 +138,20 @@ TEST("require that tensor map works") {
     EvalCtx ctx(SimpleTensorEngine::ref());
     ctx.add_tensor(ctx.make_tensor_map_input(), 1);
     auto expect = ctx.make_tensor_map_output();
-    auto fun = map(42, inject(ValueType::from_spec("tensor(x{},y{})"), 1));
+    auto fun = map(inject(ValueType::from_spec("tensor(x{},y{})"), 1), operation::Neg::f);
     EXPECT_EQUAL(ctx.type(*expect), fun->result_type);
     auto prog = ctx.compile(std::move(fun));
     TEST_DO(verify_equal(*expect, ctx.eval(*prog)));
 }
 
-TEST("require that tensor apply works") {
+TEST("require that tensor join works") {
     EvalCtx ctx(SimpleTensorEngine::ref());
     ctx.add_tensor(ctx.make_tensor_apply_lhs(), 1);
     ctx.add_tensor(ctx.make_tensor_apply_rhs(), 2);
     auto expect = ctx.make_tensor_apply_output();
-    auto fun = apply(operation::Mul(),
-                     inject(ValueType::from_spec("tensor(x{},y{})"), 1),
-                     inject(ValueType::from_spec("tensor(y{},z{})"), 2));
+    auto fun = join(inject(ValueType::from_spec("tensor(x{},y{})"), 1),
+                    inject(ValueType::from_spec("tensor(y{},z{})"), 2),
+                    operation::Mul::f);
     EXPECT_EQUAL(ctx.type(*expect), fun->result_type);
     auto prog = ctx.compile(std::move(fun));
     TEST_DO(verify_equal(*expect, ctx.eval(*prog)));
diff --git a/eval/src/tests/tensor/dense_dot_product_function/dense_dot_product_function_test.cpp b/eval/src/tests/tensor/dense_dot_product_function/dense_dot_product_function_test.cpp
index 40dfcc5b24e..17df3d21d0c 100644
--- a/eval/src/tests/tensor/dense_dot_product_function/dense_dot_product_function_test.cpp
+++ b/eval/src/tests/tensor/dense_dot_product_function/dense_dot_product_function_test.cpp
@@ -78,9 +78,6 @@ public:
             return _rhsValue;
         }
     }
-    virtual const UnaryOperation &get_map_operation(size_t) const override {
-        abort();
-    }
     double expectedDotProduct() const {
         return calcDotProduct(_lhsDenseTensor, _rhsDenseTensor);
     }
diff --git a/eval/src/tests/tensor/dense_tensor_function_compiler/dense_tensor_function_compiler_test.cpp b/eval/src/tests/tensor/dense_tensor_function_compiler/dense_tensor_function_compiler_test.cpp
index 4f6991e6418..6dcfc0791e7 100644
--- a/eval/src/tests/tensor/dense_tensor_function_compiler/dense_tensor_function_compiler_test.cpp
+++ b/eval/src/tests/tensor/dense_tensor_function_compiler/dense_tensor_function_compiler_test.cpp
@@ -16,10 +16,10 @@ TensorFunction::UP
 compileDotProduct(const vespalib::string &lhsType,
                   const vespalib::string &rhsType)
 {
-    Node_UP reduceNode = reduce(apply(Mul(),
-                                      inject(ValueType::from_spec(lhsType), 1),
-                                      inject(ValueType::from_spec(rhsType), 3)),
-                                Add(), {});
+    Node_UP reduceNode = reduce(join(inject(ValueType::from_spec(lhsType), 1),
+                                     inject(ValueType::from_spec(rhsType), 3),
+                                     Mul::f),
+                                Aggr::SUM, {});
     return DenseTensorFunctionCompiler::compile(std::move(reduceNode));
 }
 
diff --git a/eval/src/vespa/eval/eval/interpreted_function.cpp b/eval/src/vespa/eval/eval/interpreted_function.cpp
index 053dbf0288c..10fc60dd27d 100644
--- a/eval/src/vespa/eval/eval/interpreted_function.cpp
+++ b/eval/src/vespa/eval/eval/interpreted_function.cpp
@@ -177,9 +177,6 @@ struct ArgArgInput : TensorFunction::Input {
         }
         return undef_cref<Value>();
     }
-    const UnaryOperation &get_map_operation(size_t) const override {
-        return undef_cref<UnaryOperation>();
-    }
 };
 
 void op_tensor_function_arg_arg(State &state, uint64_t param) {
@@ -293,9 +290,10 @@ struct ProgramBuilder : public NodeVisitor, public NodeTraverser {
             }
             auto a = as<Symbol>(node.get_child(0).get_child(0));
             auto b = as<Symbol>(node.get_child(0).get_child(1));
-            auto ir = tensor_function::reduce(tensor_function::apply(operation::Mul(),
+            auto ir = tensor_function::reduce(tensor_function::join(
                             tensor_function::inject(types.get_type(*a), 0),
-                            tensor_function::inject(types.get_type(*b), 1)), operation::Add(), dim_list);
+                            tensor_function::inject(types.get_type(*b), 1),
+                            operation::Mul::f), Aggr::SUM, dim_list);
             auto fun = tensor_engine.compile(std::move(ir));
             const auto &meta = stash.create<TensorFunctionArgArgMeta>(std::move(fun), a->id(), b->id());
             program.emplace_back(op_tensor_function_arg_arg, wrap_param<TensorFunctionArgArgMeta>(meta));
diff --git a/eval/src/vespa/eval/eval/operation.cpp b/eval/src/vespa/eval/eval/operation.cpp
index 688bcef1d0d..16514cafaf4 100644
--- a/eval/src/vespa/eval/eval/operation.cpp
+++ b/eval/src/vespa/eval/eval/operation.cpp
@@ -73,6 +73,10 @@ template <typename T> double Op1<T>::eval(double a) const {
     return T::f(a);
 }
 
+template <typename T> Operation::op1_fun_t Op1<T>::get_f() const {
+    return T::f;
+}
+
 template <typename T> void Op2<T>::accept(OperationVisitor &visitor) const {
     visitor.visit(static_cast<const T&>(*this));
 }
@@ -85,6 +89,10 @@ template <typename T> double Op2<T>::eval(double a, double b) const {
     return T::f(a, b);
 }
 
+template <typename T> Operation::op2_fun_t Op2<T>::get_f() const {
+    return T::f;
+}
+
 namespace operation {
 double Neg::f(double a) { return -a; }
 double Not::f(double a) { return (a != 0.0) ? 0.0 : 1.0; }
diff --git a/eval/src/vespa/eval/eval/operation.h b/eval/src/vespa/eval/eval/operation.h
index 456790f9180..bd81f17aae9 100644
--- a/eval/src/vespa/eval/eval/operation.h
+++ b/eval/src/vespa/eval/eval/operation.h
@@ -20,6 +20,8 @@ struct OperationVisitor;
  * separated by the number of values they operate on.
  **/
 struct Operation {
+    using op1_fun_t = double (*)(double);
+    using op2_fun_t = double (*)(double, double);
     virtual void accept(OperationVisitor &visitor) const = 0;
     virtual ~Operation() {}
 };
@@ -39,6 +41,7 @@ const T *as(const Operation &op) { return dynamic_cast<const T *>(&op); }
 struct UnaryOperation : Operation {
     const Value &perform(const Value &a, Stash &stash) const;
     virtual double eval(double a) const = 0;
+    virtual op1_fun_t get_f() const = 0;
 };
 
 /**
@@ -48,6 +51,7 @@ struct BinaryOperation : Operation {
     const Value &perform(const Value &a, const Value &b, Stash &stash) const;
     virtual double eval(double a, double b) const = 0;
     virtual std::unique_ptr<BinaryOperation> clone() const = 0;
+    virtual op2_fun_t get_f() const = 0;
 };
 
 //-----------------------------------------------------------------------------
@@ -92,6 +96,7 @@ template <typename T>
 struct Op1 : UnaryOperation {
     virtual void accept(OperationVisitor &visitor) const override;
     virtual double eval(double a) const override;
+    virtual op1_fun_t get_f() const override;
 };
 
 template <typename T>
@@ -99,6 +104,7 @@ struct Op2 : BinaryOperation {
     virtual void accept(OperationVisitor &visitor) const override;
     virtual std::unique_ptr<BinaryOperation> clone() const override;
     virtual double eval(double a, double b) const final override;
+    virtual op2_fun_t get_f() const override;
 };
 
 //-----------------------------------------------------------------------------
@@ -108,7 +114,7 @@ struct Op2 : BinaryOperation {
  * and lambdas.
  **/
 struct CustomUnaryOperation : Op1<CustomUnaryOperation> {
-    static double f(double);
+    static double f(double) { return error_value; }
 };
 
 //-----------------------------------------------------------------------------
diff --git a/eval/src/vespa/eval/eval/tensor_function.cpp b/eval/src/vespa/eval/eval/tensor_function.cpp
index afcd157baa8..0dcc930087f 100644
--- a/eval/src/vespa/eval/eval/tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/tensor_function.cpp
@@ -13,7 +13,7 @@ namespace tensor_function {
 void Inject::accept(TensorFunctionVisitor &visitor) const { visitor.visit(*this); }
 void Reduce::accept(TensorFunctionVisitor &visitor) const { visitor.visit(*this); }
 void Map   ::accept(TensorFunctionVisitor &visitor) const { visitor.visit(*this); }
-void Apply ::accept(TensorFunctionVisitor &visitor) const { visitor.visit(*this); }
+void Join  ::accept(TensorFunctionVisitor &visitor) const { visitor.visit(*this); }
 
 //-----------------------------------------------------------------------------
 
@@ -26,26 +26,26 @@ Inject::eval(const Input &input, Stash &) const
 const Value &
 Reduce::eval(const Input &input, Stash &stash) const 
 {
-    const Tensor &a = *tensor->eval(input, stash).as_tensor();
-    const TensorEngine &engine = a.engine();
-    return engine.reduce(a, *op, dimensions, stash);
+    const Value &a = tensor->eval(input, stash);
+    const TensorEngine &engine = a.as_tensor()->engine();
+    return engine.reduce(a, aggr, dimensions, stash);
 }
 
 const Value &
 Map::eval(const Input &input, Stash &stash) const
 {
-    const Tensor &a = *tensor->eval(input, stash).as_tensor();
-    const TensorEngine &engine = a.engine();
-    return engine.map(input.get_map_operation(map_operation_id), a, stash);
+    const Value &a = tensor->eval(input, stash);
+    const TensorEngine &engine = a.as_tensor()->engine();
+    return engine.map(a, function, stash);
 }
 
 const Value &
-Apply::eval(const Input &input, Stash &stash) const
+Join::eval(const Input &input, Stash &stash) const
 {
-    const Tensor &a = *lhs_tensor->eval(input, stash).as_tensor();
-    const Tensor &b = *rhs_tensor->eval(input, stash).as_tensor();
-    const TensorEngine &engine = a.engine();
-    return engine.apply(*op, a, b, stash);
+    const Value &a = lhs_tensor->eval(input, stash);
+    const Value &b = rhs_tensor->eval(input, stash);
+    const TensorEngine &engine = a.as_tensor()->engine();
+    return engine.join(a, b, function, stash);
 }
 
 //-----------------------------------------------------------------------------
@@ -54,19 +54,19 @@ Node_UP inject(const ValueType &type, size_t tensor_id) {
     return std::make_unique<Inject>(type, tensor_id);
 }
 
-Node_UP reduce(Node_UP tensor, const BinaryOperation &op, const std::vector<vespalib::string> &dimensions) {
+Node_UP reduce(Node_UP tensor, Aggr aggr, const std::vector<vespalib::string> &dimensions) {
     ValueType result_type = tensor->result_type.reduce(dimensions);
-    return std::make_unique<Reduce>(result_type, std::move(tensor), op.clone(), dimensions);
+    return std::make_unique<Reduce>(result_type, std::move(tensor), aggr, dimensions);
 }
 
-Node_UP map(size_t map_operation_id, Node_UP tensor) {
+Node_UP map(Node_UP tensor, map_fun_t function) {
     ValueType result_type = tensor->result_type;
-    return std::make_unique<Map>(result_type, map_operation_id, std::move(tensor));
+    return std::make_unique<Map>(result_type, std::move(tensor), function);
 }
 
-Node_UP apply(const BinaryOperation &op, Node_UP lhs_tensor, Node_UP rhs_tensor) {
+Node_UP join(Node_UP lhs_tensor, Node_UP rhs_tensor, join_fun_t function) {
     ValueType result_type = ValueType::join(lhs_tensor->result_type, rhs_tensor->result_type);
-    return std::make_unique<Apply>(result_type, op.clone(), std::move(lhs_tensor), std::move(rhs_tensor));
+    return std::make_unique<Join>(result_type, std::move(lhs_tensor), std::move(rhs_tensor), function);
 }
 
 } // namespace vespalib::eval::tensor_function
diff --git a/eval/src/vespa/eval/eval/tensor_function.h b/eval/src/vespa/eval/eval/tensor_function.h
index e3945915e37..cff21b7b9aa 100644
--- a/eval/src/vespa/eval/eval/tensor_function.h
+++ b/eval/src/vespa/eval/eval/tensor_function.h
@@ -7,6 +7,7 @@
 #include <vespa/vespalib/stllike/string.h>
 #include "value_type.h"
 #include "operation.h"
+#include "aggr.h"
 
 namespace vespalib {
 
@@ -35,7 +36,6 @@ struct TensorFunction
      **/
     struct Input {
         virtual const Value &get_tensor(size_t id) const = 0;
-        virtual const UnaryOperation &get_map_operation(size_t id) const = 0;
         virtual ~Input() {}
     };
 
@@ -60,6 +60,9 @@ struct TensorFunctionVisitor;
 
 namespace tensor_function {
 
+using map_fun_t = double (*)(double);
+using join_fun_t = double (*)(double, double);
+
 /**
  * Interface used to describe a tensor function as a tree of nodes
  * with information about operation sequencing and intermediate result
@@ -102,46 +105,46 @@ struct Inject : Node {
 
 struct Reduce : Node {
     Node_UP tensor;
-    std::unique_ptr<BinaryOperation> op;
+    Aggr aggr;
     std::vector<vespalib::string> dimensions;
     Reduce(const ValueType &result_type_in,
            Node_UP tensor_in,
-           std::unique_ptr<BinaryOperation> op_in,
+           Aggr aggr_in,
            const std::vector<vespalib::string> &dimensions_in)
-        : Node(result_type_in), tensor(std::move(tensor_in)), op(std::move(op_in)), dimensions(dimensions_in) {}
+        : Node(result_type_in), tensor(std::move(tensor_in)), aggr(aggr_in), dimensions(dimensions_in) {}
     void accept(TensorFunctionVisitor &visitor) const override;
     const Value &eval(const Input &input, Stash &stash) const override;
 };
 
 struct Map : Node {
-    size_t map_operation_id;
     Node_UP tensor;
+    map_fun_t function;    
     Map(const ValueType &result_type_in,
-        size_t map_operation_id_in,
-        Node_UP tensor_in)
-        : Node(result_type_in), map_operation_id(map_operation_id_in), tensor(std::move(tensor_in)) {}
+        Node_UP tensor_in,
+        map_fun_t function_in)
+        : Node(result_type_in), tensor(std::move(tensor_in)), function(function_in) {}
     void accept(TensorFunctionVisitor &visitor) const override;
     const Value &eval(const Input &input, Stash &stash) const override;
 };
 
-struct Apply : Node {
-    std::unique_ptr<BinaryOperation> op;
+struct Join : Node {
     Node_UP lhs_tensor;
     Node_UP rhs_tensor;
-    Apply(const ValueType &result_type_in,
-          std::unique_ptr<BinaryOperation> op_in,
-          Node_UP lhs_tensor_in,
-          Node_UP rhs_tensor_in)
-        : Node(result_type_in), op(std::move(op_in)),
-          lhs_tensor(std::move(lhs_tensor_in)), rhs_tensor(std::move(rhs_tensor_in)) {}
+    join_fun_t function;    
+    Join(const ValueType &result_type_in,
+         Node_UP lhs_tensor_in,
+         Node_UP rhs_tensor_in,
+         join_fun_t function_in)
+        : Node(result_type_in), lhs_tensor(std::move(lhs_tensor_in)),
+          rhs_tensor(std::move(rhs_tensor_in)), function(function_in) {}
     void accept(TensorFunctionVisitor &visitor) const override;
     const Value &eval(const Input &input, Stash &stash) const override;
 };
 
 Node_UP inject(const ValueType &type, size_t tensor_id);
-Node_UP reduce(Node_UP tensor, const BinaryOperation &op, const std::vector<vespalib::string> &dimensions);
-Node_UP map(size_t map_operation_id, Node_UP tensor);
-Node_UP apply(const BinaryOperation &op, Node_UP lhs_tensor, Node_UP rhs_tensor);
+Node_UP reduce(Node_UP tensor, Aggr aggr, const std::vector<vespalib::string> &dimensions);
+Node_UP map(Node_UP tensor, map_fun_t function);
+Node_UP join(Node_UP lhs_tensor, Node_UP rhs_tensor, join_fun_t function);
 
 } // namespace vespalib::eval::tensor_function
 
@@ -149,7 +152,7 @@ struct TensorFunctionVisitor {
     virtual void visit(const tensor_function::Inject &) = 0;
     virtual void visit(const tensor_function::Reduce &) = 0;
     virtual void visit(const tensor_function::Map    &) = 0;
-    virtual void visit(const tensor_function::Apply  &) = 0;
+    virtual void visit(const tensor_function::Join   &) = 0;
     virtual ~TensorFunctionVisitor() {}
 };
 
diff --git a/eval/src/vespa/eval/eval/test/tensor_conformance.cpp b/eval/src/vespa/eval/eval/test/tensor_conformance.cpp
index 828c01b47f2..ebfd3c11634 100644
--- a/eval/src/vespa/eval/eval/test/tensor_conformance.cpp
+++ b/eval/src/vespa/eval/eval/test/tensor_conformance.cpp
@@ -23,6 +23,9 @@ using slime::Cursor;
 using slime::Inspector;
 using slime::JsonFormat;
 
+using map_fun_t = TensorEngine::map_fun_t;
+using join_fun_t = TensorEngine::join_fun_t;
+
 double as_double(const TensorSpec &spec) {
     return spec.cells().empty() ? 0.0 : spec.cells().begin()->second.value;
 }
@@ -273,19 +276,14 @@ struct ImmediateRename : Eval {
 
 const size_t tensor_id_a = 11;
 const size_t tensor_id_b = 12;
-const size_t map_operation_id = 22;
 
 // input used when evaluating in retained mode
 struct Input : TensorFunction::Input {
     std::vector<TensorValue> tensors;
-    const UnaryOperation *map_op;
-    Input(std::unique_ptr<Tensor> a) : tensors(), map_op(nullptr) {
-        tensors.emplace_back(std::move(a));
-    }
-    Input(std::unique_ptr<Tensor> a, const UnaryOperation &op) : tensors(), map_op(&op) {
+    Input(std::unique_ptr<Tensor> a) : tensors() {
         tensors.emplace_back(std::move(a));
     }
-    Input(std::unique_ptr<Tensor> a, std::unique_ptr<Tensor> b) : tensors(), map_op(nullptr) {
+    Input(std::unique_ptr<Tensor> a, std::unique_ptr<Tensor> b) : tensors() {
         tensors.emplace_back(std::move(a));
         tensors.emplace_back(std::move(b));
     }
@@ -294,23 +292,18 @@ struct Input : TensorFunction::Input {
         ASSERT_GREATER(tensors.size(), offset);
         return tensors[offset];
     }
-    const UnaryOperation &get_map_operation(size_t id) const override {
-        ASSERT_TRUE(map_op != nullptr);
-        ASSERT_EQUAL(id, map_operation_id);
-        return *map_op;
-    }
 };
 
 // evaluate tensor reduce operation using tensor engine retained api
 struct RetainedReduce : Eval {
-    const BinaryOperation &op;
+    Aggr aggr;
     std::vector<vespalib::string> dimensions;
-    RetainedReduce(const BinaryOperation &op_in) : op(op_in), dimensions() {}
-    RetainedReduce(const BinaryOperation &op_in, const vespalib::string &dimension)
-        : op(op_in), dimensions({dimension}) {}
+    RetainedReduce(Aggr aggr_in) : aggr(aggr_in), dimensions() {}
+    RetainedReduce(Aggr aggr_in, const vespalib::string &dimension)
+        : aggr(aggr_in), dimensions({dimension}) {}
     Result eval(const TensorEngine &engine, const TensorSpec &a) const override {
         auto a_type = ValueType::from_spec(a.type());
-        auto ir = tensor_function::reduce(tensor_function::inject(a_type, tensor_id_a), op, dimensions);
+        auto ir = tensor_function::reduce(tensor_function::inject(a_type, tensor_id_a), aggr, dimensions);
         ValueType expect_type = ir->result_type;
         auto fun = engine.compile(std::move(ir));
         Input input(engine.create(a));
@@ -321,28 +314,29 @@ struct RetainedReduce : Eval {
 
 // evaluate tensor map operation using tensor engine retained api
 struct RetainedMap : Eval {
-    const UnaryOperation &op;
-    RetainedMap(const UnaryOperation &op_in) : op(op_in) {}
+    map_fun_t function;
+    RetainedMap(map_fun_t function_in) : function(function_in) {}
     Result eval(const TensorEngine &engine, const TensorSpec &a) const override {
         auto a_type = ValueType::from_spec(a.type());
-        auto ir = tensor_function::map(map_operation_id, tensor_function::inject(a_type, tensor_id_a));
+        auto ir = tensor_function::map(tensor_function::inject(a_type, tensor_id_a), function);
         ValueType expect_type = ir->result_type;
         auto fun = engine.compile(std::move(ir));
-        Input input(engine.create(a), op);
+        Input input(engine.create(a));
         Stash stash;
         return Result(check_type(fun->eval(input, stash), expect_type));
     }
 };
 
-// evaluate tensor apply operation using tensor engine retained api
-struct RetainedApply : Eval {
-    const BinaryOperation &op;
-    RetainedApply(const BinaryOperation &op_in) : op(op_in) {}
+// evaluate tensor join operation using tensor engine retained api
+struct RetainedJoin : Eval {
+    join_fun_t function;
+    RetainedJoin(join_fun_t function_in) : function(function_in) {}
     Result eval(const TensorEngine &engine, const TensorSpec &a, const TensorSpec &b) const override {
         auto a_type = ValueType::from_spec(a.type());
         auto b_type = ValueType::from_spec(b.type());
-        auto ir = tensor_function::apply(op, tensor_function::inject(a_type, tensor_id_a),
-                                         tensor_function::inject(b_type, tensor_id_b));
+        auto ir = tensor_function::join(tensor_function::inject(a_type, tensor_id_a),
+                                        tensor_function::inject(b_type, tensor_id_b),
+                                        function);
         ValueType expect_type = ir->result_type;
         auto fun = engine.compile(std::move(ir));
         Input input(engine.create(a), engine.create(b));
@@ -519,7 +513,7 @@ struct TestContext {
                 TEST_DO(verify_reduce_result(Expr_T(expr), input, expect));
                 TEST_DO(verify_reduce_result(ImmediateReduceOld(op, domain.dimension), input, expect));
                 TEST_DO(verify_reduce_result(ImmediateReduce(aggr, domain.dimension), input, expect));
-                TEST_DO(verify_reduce_result(RetainedReduce(op, domain.dimension), input, expect));
+                TEST_DO(verify_reduce_result(RetainedReduce(aggr, domain.dimension), input, expect));
             }
             {
                 Eval::Result expect = ImmediateReduceOld(op).eval(ref_engine, input);
@@ -530,7 +524,7 @@ struct TestContext {
                 TEST_DO(verify_reduce_result(Expr_T(expr), input, expect));
                 TEST_DO(verify_reduce_result(ImmediateReduceOld(op), input, expect));
                 TEST_DO(verify_reduce_result(ImmediateReduce(aggr), input, expect));
-                TEST_DO(verify_reduce_result(RetainedReduce(op), input, expect));
+                TEST_DO(verify_reduce_result(RetainedReduce(aggr), input, expect));
             }
         }
     }
@@ -556,6 +550,7 @@ struct TestContext {
                         AggrNames::name_of(aggr)->c_str(), domain.dimension.c_str());
                 TEST_DO(verify_reduce_result(Expr_T(expr), input, expect));
                 TEST_DO(verify_reduce_result(ImmediateReduce(aggr, domain.dimension), input, expect));
+                TEST_DO(verify_reduce_result(RetainedReduce(aggr, domain.dimension), input, expect));
             }
             {
                 Eval::Result expect = ImmediateReduce(aggr).eval(ref_engine, input);
@@ -564,6 +559,7 @@ struct TestContext {
                 vespalib::string expr = make_string("reduce(a,%s)", AggrNames::name_of(aggr)->c_str());
                 TEST_DO(verify_reduce_result(Expr_T(expr), input, expect));
                 TEST_DO(verify_reduce_result(ImmediateReduce(aggr), input, expect));
+                TEST_DO(verify_reduce_result(RetainedReduce(aggr), input, expect));
             }
         }
     }
@@ -598,8 +594,8 @@ struct TestContext {
 
     void test_map_op(const vespalib::string &expr, const UnaryOperation &op, const Sequence &seq) {
         TEST_DO(test_map_op(ImmediateMapOld(op), op, seq));
-        TEST_DO(test_map_op(ImmediateMap(UnaryOperationProxy(op)), op, seq));
-        TEST_DO(test_map_op(RetainedMap(op), op, seq));
+        TEST_DO(test_map_op(ImmediateMap(op.get_f()), op, seq));
+        TEST_DO(test_map_op(RetainedMap(op.get_f()), op, seq));
         TEST_DO(test_map_op(Expr_T(expr), op, seq));
         TEST_DO(test_map_op(Expr_T(make_string("map(x,f(a)(%s))", expr.c_str())), op, seq));
     }
@@ -874,8 +870,8 @@ struct TestContext {
 
     void test_apply_op(const vespalib::string &expr, const BinaryOperation &op, const Sequence &seq) {
         TEST_DO(test_apply_op(ImmediateApplyOld(op), op, seq));
-        TEST_DO(test_apply_op(ImmediateJoin(BinaryOperationProxy(op)), op, seq));
-        TEST_DO(test_apply_op(RetainedApply(op), op, seq));
+        TEST_DO(test_apply_op(ImmediateJoin(op.get_f()), op, seq));
+        TEST_DO(test_apply_op(RetainedJoin(op.get_f()), op, seq));
         TEST_DO(test_apply_op(Expr_TT(expr), op, seq));
         TEST_DO(test_apply_op(Expr_TT(make_string("join(x,y,f(a,b)(%s))", expr.c_str())), op, seq));
     }
diff --git a/eval/src/vespa/eval/eval/test/tensor_model.hpp b/eval/src/vespa/eval/eval/test/tensor_model.hpp
index 8f031174289..53a4f85a5f4 100644
--- a/eval/src/vespa/eval/eval/test/tensor_model.hpp
+++ b/eval/src/vespa/eval/eval/test/tensor_model.hpp
@@ -106,7 +106,11 @@ struct Mask2Seq : Sequence {
 
 // custom op1
 struct MyOp : CustomUnaryOperation {
-    double eval(double a) const override { return ((a + 1) * 2); }
+    static double f(double a) {
+        return ((a + 1) * 2);
+    }
+    double eval(double a) const override { return f(a); }
+    op1_fun_t get_f() const override { return f; }
 };
 
 // A collection of labels for a single dimension
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_function_compiler.cpp b/eval/src/vespa/eval/tensor/dense/dense_tensor_function_compiler.cpp
index a82e1f282ba..ac7b5d79a11 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_function_compiler.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_function_compiler.cpp
@@ -16,13 +16,6 @@ namespace tensor {
 
 namespace {
 
-template <typename T>
-bool    
-isType(const BinaryOperation &op)
-{
-    return (as<T>(op) != nullptr);
-}
-
 bool
 willReduceAllDimensions(const std::vector<vespalib::string> &dimensions)
 {
@@ -47,11 +40,11 @@ struct DotProductFunctionCompiler
 {
     static TensorFunction::UP compile(Node_UP expr) {
         const Reduce *reduce = as<Reduce>(*expr);
-        if (reduce && isType<Add>(*reduce->op) && willReduceAllDimensions(reduce->dimensions)) {
-            const Apply *apply = as<Apply>(*reduce->tensor);
-            if (apply && isType<Mul>(*apply->op)) {
-                const Inject *lhsTensor = as<Inject>(*apply->lhs_tensor);
-                const Inject *rhsTensor = as<Inject>(*apply->rhs_tensor);
+        if (reduce && (reduce->aggr == Aggr::SUM) && willReduceAllDimensions(reduce->dimensions)) {
+            const Join *join = as<Join>(*reduce->tensor);
+            if (join && (join->function == Mul::f)) {
+                const Inject *lhsTensor = as<Inject>(*join->lhs_tensor);
+                const Inject *rhsTensor = as<Inject>(*join->rhs_tensor);
                 if (lhsTensor && rhsTensor &&
                     isCompatibleTensorsForDotProduct(lhsTensor->result_type, rhsTensor->result_type))
                 {