use withTransformedExpressions for wiring

1 files changed, 5 insertions, 14 deletions

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 bc71e51655c..8cdd5387acd 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
@@ -96,6 +96,11 @@ class BindingExtractor {
             for (ExpressionNode child : tfn.children()) {
                 result.merge(extractBindTargets(child));
             }
+            // ignore return value:
+            tfn.withTransformedExpressions(expr -> {
+                    result.merge(extractBindTargets(expr));
+                    return expr;
+                });
             var f = tfn.function();
             if (f instanceof Generate) {
                 var tt = f.type(null);
@@ -103,20 +108,6 @@ class BindingExtractor {
                     result.removeTarget(dim.name());
                 }
             }
-            else if (f instanceof Slice<?> s) {
-                for (var selectorFunc : s.selectorFunctions()) {
-                    if (selectorFunc instanceof TensorFunctionNode.ExpressionTensorFunction expr) {
-                        result.merge(extractBindTargets(expr.wrappedExpression()));
-                    }
-                }
-            }
-            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)) {