1 files changed, 17 insertions, 0 deletions
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Slice.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Slice.java
index 87e24306031..066d75bcd9c 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Slice.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Slice.java
@@ -16,6 +16,7 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
@@ -56,6 +57,22 @@ public class Slice<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETY
         return result;
     }
 
+    public TensorFunction<NAMETYPE> withTransformedFunctions(
+            Function<ScalarFunction<NAMETYPE>, ScalarFunction<NAMETYPE>> transformer)
+    {
+        List<DimensionValue<NAMETYPE>> transformedAddress = new ArrayList<>();
+        for (var orig : subspaceAddress) {
+            var idxFun = orig.index();
+            if (idxFun.isPresent()) {
+                var transformed = transformer.apply(idxFun.get());
+                transformedAddress.add(new DimensionValue<NAMETYPE>(orig.dimension(), transformed));
+            } else {
+                transformedAddress.add(orig);
+            }
+        }
+        return new Slice<>(argument, transformedAddress);
+    }
+
     @Override
     public Slice<NAMETYPE> withArguments(List<TensorFunction<NAMETYPE>> arguments) {
         if (arguments.size() != 1)