Repository: vespa-engine/vespa
File: vespajlib/src/main/java/com/yahoo/tensor/TensorDataSource.java (Java)

Code before:

    // Hex encoding methods adapted from JsonFormat
    private static final byte[] hexDigits = "0123456789ABCDEF".getBytes(StandardCharsets.US_ASCII);

    private static byte[] asHexString(IndexedTensor tensor) {
        return asHexString(tensor.sizeAsInt(), tensor.type().valueType(),
                           i -> tensor.get(i), i -> tensor.getFloat(i));
    }

    private static byte[] asHexString(int denseSize, TensorType.Value cellType,
                                      java.util.function.Function<Integer, Double> dblSrc,
                                      java.util.function.Function<Integer, Float> fltSrc)
    {
        int nibblesPerCell = switch (cellType) {
            case DOUBLE   -> 16;
            case FLOAT    -> 8;
            case BFLOAT16 -> 4;
            case INT8     -> 2;
        };
        byte[] result = new byte[nibblesPerCell * denseSize];
        int idx = 0;
        for (int i = 0; i < denseSize; i++) {
            long bits = switch (cellType) {
                case DOUBLE   -> Double.doubleToRawLongBits(dblSrc.apply(i));
                case FLOAT    -> Float.floatToRawIntBits(fltSrc.apply(i));
                case BFLOAT16 -> Float.floatToRawIntBits(fltSrc.apply(i));
                case INT8     -> fltSrc.apply(i).byteValue();
            };
            for (int nibble = nibblesPerCell; nibble-- > 0; ) {
                int digit = (int) (bits >> (4 * nibble)) & 0xF;
                result[idx++] = hexDigits[digit];
            }
        }
        if (idx != result.length) {
            throw new IllegalStateException("Did not fill result["+result.length+"], final idx=" + idx);
        }
        return result;
    }

    private void wrapStart(DataSink sink) {
        if (wrapAndType) {
            ensureObject(sink);
        }
    }
    private void ensureObjectEnded(DataSink sink) {
        if (inObject) {
            sink.endObject();
            inObject = false;
        }
    }
    private void startField(String fieldName, DataSink sink) {

Review comment (suggestion):
```suggestion
                case BFLOAT16 -> Float.floatToRawIntBits(fltSrc.apply(i)) >>> 16;
```

Code after:

    // Hex encoding methods adapted from JsonFormat
    private static final byte[] hexDigits = "0123456789ABCDEF".getBytes(StandardCharsets.US_ASCII);

    private static byte[] asHexString(IndexedTensor tensor) {
        return asHexString(tensor.sizeAsInt(), tensor.type().valueType(),
                           i -> tensor.get(i), i -> tensor.getFloat(i));
    }

    private static byte[] asHexString(int denseSize, TensorType.Value cellType,
                                      java.util.function.Function<Integer, Double> dblSrc,
                                      java.util.function.Function<Integer, Float> fltSrc)
    {
        int nibblesPerCell = switch (cellType) {
            case DOUBLE   -> 16;
            case FLOAT    -> 8;
            case BFLOAT16 -> 4;
            case INT8     -> 2;
        };
        byte[] result = new byte[nibblesPerCell * denseSize];
        int idx = 0;
        for (int i = 0; i < denseSize; i++) {
            long bits = switch (cellType) {
                case DOUBLE   -> Double.doubleToRawLongBits(dblSrc.apply(i));
                case FLOAT    -> Float.floatToRawIntBits(fltSrc.apply(i));
                case BFLOAT16 -> Float.floatToRawIntBits(fltSrc.apply(i)) >>> 16;
                case INT8     -> fltSrc.apply(i).byteValue();
            };
            for (int nibble = nibblesPerCell; nibble-- > 0; ) {
                int digit = (int) (bits >> (4 * nibble)) & 0xF;
                result[idx++] = hexDigits[digit];
            }
        }
        if (idx != result.length) {
            throw new IllegalStateException("Did not fill result["+result.length+"], final idx=" + idx);
        }
        return result;
    }

    private void wrapStart(DataSink sink) {
        if (wrapAndType) {
            ensureObject(sink);
        }
    }
    private void ensureObjectEnded(DataSink sink) {
        if (inObject) {
            sink.endObject();
            inObject = false;
        }
    }
    private void startField(String fieldName, DataSink sink) {