Add bfloat16 and int8 to Java cell_cast function

1 files changed, 7 insertions, 3 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/CellCast.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/CellCast.java
index d052e383c85..d853c4a9069 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/CellCast.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/CellCast.java
@@ -64,10 +64,14 @@ public class CellCast<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAM
         TensorType.Value fromValueType = tensor.type().valueType();
         for (Iterator<Tensor.Cell> i = tensor.cellIterator(); i.hasNext(); ) {
             Tensor.Cell cell = i.next();
-            if (fromValueType == TensorType.Value.FLOAT) {
-                builder.cell(cell.getKey(), cell.getFloatValue());
-            } else if (fromValueType == TensorType.Value.DOUBLE) {
+            if (fromValueType == TensorType.Value.DOUBLE) {
                 builder.cell(cell.getKey(), cell.getDoubleValue());
+            } else if (fromValueType == TensorType.Value.FLOAT) {
+                builder.cell(cell.getKey(), cell.getFloatValue());
+            } else if (fromValueType == TensorType.Value.BFLOAT16) {
+                builder.cell(cell.getKey(), cell.getFloatValue());
+            } else if (fromValueType == TensorType.Value.INT8) {
+                builder.cell(cell.getKey(), cell.getFloatValue());
             } else {
                 builder.cell(cell.getKey(), cell.getValue());
             }