From 89862f601ce193c1b8aa96d9979ad5ef57c1a505 Mon Sep 17 00:00:00 2001 From: Arne Juul Date: Fri, 17 Mar 2023 12:16:56 +0000 Subject: handle DynamicTensor --- .../main/java/ai/vespa/models/evaluation/BindingExtractor.java | 8 ++++++++ 1 file changed, 8 insertions(+) (limited to 'model-evaluation') diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/BindingExtractor.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/BindingExtractor.java index 126e9f9f4e6..bc71e51655c 100644 --- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/BindingExtractor.java +++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/BindingExtractor.java @@ -8,6 +8,7 @@ import com.yahoo.searchlib.rankingexpression.rule.ConstantNode; import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode; import com.yahoo.searchlib.rankingexpression.rule.ReferenceNode; import com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode; +import com.yahoo.tensor.functions.DynamicTensor; import com.yahoo.tensor.functions.Generate; import com.yahoo.tensor.functions.Slice; @@ -109,6 +110,13 @@ class BindingExtractor { } } } + else if (f instanceof DynamicTensor d) { + for (var tf : d.cellGeneratorFunctions()) { + if (tf instanceof TensorFunctionNode.ExpressionTensorFunction expr) { + result.merge(extractBindTargets(expr.wrappedExpression())); + } + } + } return result; } else if (isOnnx(node)) { -- cgit v1.2.3