Propagate type information

1 files changed, 7 insertions, 4 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
index 174a8e4c435..01c681bfb36 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
@@ -70,15 +70,13 @@ public class Join extends PrimitiveTensorFunction {
         return typeBuilder.build();
     }
 
-    public TensorFunction argumentA() { return argumentA; }
-    public TensorFunction argumentB() { return argumentB; }
     public DoubleBinaryOperator combinator() { return combinator; }
 
     @Override
-    public List<TensorFunction> functionArguments() { return ImmutableList.of(argumentA, argumentB); }
+    public List<TensorFunction> arguments() { return ImmutableList.of(argumentA, argumentB); }
 
     @Override
-    public TensorFunction replaceArguments(List<TensorFunction> arguments) {
+    public TensorFunction withArguments(List<TensorFunction> arguments) {
         if ( arguments.size() != 2)
             throw new IllegalArgumentException("Join must have 2 arguments, got " + arguments.size());
         return new Join(arguments.get(0), arguments.get(1), combinator);
@@ -95,6 +93,11 @@ public class Join extends PrimitiveTensorFunction {
     }
 
     @Override
+    public TensorType type(EvaluationContext context) {
+        return new TensorType.Builder(argumentA.type(context), argumentB.type(context)).build();
+    }
+
+    @Override
     public Tensor evaluate(EvaluationContext context) {
         Tensor a = argumentA.evaluate(context);
         Tensor b = argumentB.evaluate(context);