Preserve tensor type

1 files changed, 3 insertions, 3 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Rename.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Rename.java
index 4f0080ca3cc..39d555d1632 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Rename.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Rename.java
@@ -69,12 +69,12 @@ public class Rename extends PrimitiveTensorFunction {
             toIndexes[i] = renamedType.indexOfDimension(newDimensionName).get();
         }
             
-        ImmutableMap.Builder<TensorAddress, Double> renamedCells = new ImmutableMap.Builder<>();
+        Tensor.Builder builder = Tensor.Builder.of(renamedType);
         for (Map.Entry<TensorAddress, Double> cell : tensor.cells().entrySet()) {
             TensorAddress renamedAddress = rename(cell.getKey(), toIndexes);
-            renamedCells.put(renamedAddress, cell.getValue());
+            builder.cell(renamedAddress, cell.getValue());
         }
-        return new MappedTensor(renamedType, renamedCells.build());
+        return builder.build();
     }
 
     private TensorType rename(TensorType type, Map<String, String> fromToMap) {