Propagate type information

1 files changed, 3 insertions, 3 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/evaluation/VariableTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/VariableTensor.java
index 5f809a3d2b1..acb2363cba4 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/evaluation/VariableTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/evaluation/VariableTensor.java
@@ -44,15 +44,15 @@ public class VariableTensor extends PrimitiveTensorFunction {
     public PrimitiveTensorFunction toPrimitive() { return this; }
 
     @Override
-    public TensorType type(TypeContext context) {
-        TensorType givenType = context.getType(new TypeContext.Name(name));
+    public <NAMETYPE extends TypeContext.Name> TensorType type(TypeContext<NAMETYPE> context) {
+        TensorType givenType = context.getType(name);
         if (givenType == null) return null;
         verifyType(givenType);
         return givenType;
     }
 
     @Override
-    public Tensor evaluate(EvaluationContext context) {
+    public <NAMETYPE extends TypeContext.Name> Tensor evaluate(EvaluationContext<NAMETYPE> context) {
         Tensor tensor = context.getTensor(name);
         if (tensor == null) return null;
         verifyType(tensor.type());