Repository: deepjavalibrary/djl (java)

Issue:
How to run FLOAT16 OnnxRuntime models
It looks like FLOAT16 has not been supported by now
```
Caused by: java.lang.UnsupportedOperationException: type is not supported: FLOAT16
at ai.djl.onnxruntime.engine.OrtUtils.toDataType(OrtUtils.java:101)
at ai.djl.onnxruntime.engine.OrtNDArray.getDataType(OrtNDArray.java:65)
at ai.djl.onnxruntime.engine.OrtNDArray.toByteBuffer(OrtNDArray.java:121)
at ai.djl.pytorch.engine.PtNDManager.from(PtNDManager.java:55)
at ai.djl.pytorch.engine.PtNDManager.from(PtNDManager.java:31)
at ai.djl.ndarray.NDArrayAdapter.getAlternativeArray(NDArrayAdapter.java:1315)
at ai.djl.ndarray.NDArrayAdapter.split(NDArrayAdapter.java:876)
at ai.djl.ndarray.NDArray.split(NDArray.java:3173)
at ai.djl.translate.StackBatchifier.unbatchify(StackBatchifier.java:118)
at ai.djl.huggingface.translator.CrossEncoderBatchTranslator.processOutput(CrossEncoderBatchTranslator.java:60)
at ai.djl.huggingface.translator.CrossEncoderBatchTranslator.processOutput(CrossEncoderBatchTranslator.java:30)
at ai.djl.inference.Predictor.batchPredict(Predictor.java:173)
... 5 more
```

Pull request description:
[onnx] Adds fp16 and bfp16 support for OnnxRuntime
Fixes #3190

## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


Patch:
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtUtils.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtUtils.java
index 07dbfa7c7..50143cd68 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtUtils.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtUtils.java
@@ -36,30 +36,33 @@ final class OrtUtils {
         if (shape.size() == 0) {
             throw new UnsupportedOperationException("OnnxRuntime doesn't support 0 length tensor.");
         }
-        if (data instanceof ByteBuffer) {
-            data = dataType.asDataType((ByteBuffer) data);
-        }
         long[] sh = shape.getShape();
         try {
             switch (dataType) {
                 case FLOAT32:
-                    return OnnxTensor.createTensor(env, (FloatBuffer) data, sh);
+                    return OnnxTensor.createTensor(env, asFloatBuffer(data), sh);
                 case FLOAT64:
-                    return OnnxTensor.createTensor(env, (DoubleBuffer) data, sh);
+                    return OnnxTensor.createTensor(env, asDoubleBuffer(data), sh);
+                case FLOAT16:
+                    return OnnxTensor.createTensor(
+                            env, (ByteBuffer) data, sh, OnnxJavaType.FLOAT16);
+                case BFLOAT16:
+                    return OnnxTensor.createTensor(
+                            env, (ByteBuffer) data, sh, OnnxJavaType.BFLOAT16);
                 case INT32:
-                    return OnnxTensor.createTensor(env, (IntBuffer) data, sh);
+                    return OnnxTensor.createTensor(env, asIntBuffer(data), sh);
                 case INT64:
-                    return OnnxTensor.createTensor(env, (LongBuffer) data, sh);
+                    return OnnxTensor.createTensor(env, asLongBuffer(data), sh);
                 case INT8:
                     return OnnxTensor.createTensor(env, (ByteBuffer) data, sh, OnnxJavaType.INT8);
                 case UINT8:
                     return OnnxTensor.createTensor(env, (ByteBuffer) data, sh, OnnxJavaType.UINT8);
+                case BOOLEAN:
+                    return OnnxTensor.createTensor(env, (ByteBuffer) data, sh, OnnxJavaType.BOOL);
                 case STRING:
                     throw new UnsupportedOperationException(
                             "Use toTensor(OrtEnvironment env, String[] inputs, Shape shape)"
                                     + " instead.");
-                case BOOLEAN:
-                case FLOAT16:
                 default:
                     throw new UnsupportedOperationException("Data type not supported: " + dataType);
             }
@@ -81,6 +84,10 @@ final class OrtUtils {
         switch (javaType) {
             case FLOAT:
                 return DataType.FLOAT32;
+            case FLOAT16:
+                return DataType.FLOAT16;
+            case BFLOAT16:
+                return DataType.BFLOAT16;
             case DOUBLE:
                 return DataType.FLOAT64;
             case INT8:
@@ -101,4 +108,32 @@ final class OrtUtils {
                 throw new UnsupportedOperationException("type is not supported: " + javaType);
         }
     }
+
+    private static FloatBuffer asFloatBuffer(Buffer data) {
+        if (data instanceof ByteBuffer) {
+            return ((ByteBuffer) data).asFloatBuffer();
+        }
+        return (FloatBuffer) data;
+    }
+
+    private static DoubleBuffer asDoubleBuffer(Buffer data) {
+        if (data instanceof ByteBuffer) {
+            return ((ByteBuffer) data).asDoubleBuffer();
+        }
+        return (DoubleBuffer) data;
+    }
+
+    private static IntBuffer asIntBuffer(Buffer data) {
+        if (data instanceof ByteBuffer) {
+            return ((ByteBuffer) data).asIntBuffer();
+        }
+        return (IntBuffer) data;
+    }
+
+    private static LongBuffer asLongBuffer(Buffer data) {
+        if (data instanceof ByteBuffer) {
+            return ((ByteBuffer) data).asLongBuffer();
+        }
+        return (LongBuffer) data;
+    }
 }


Test patch:
diff --git a/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java b/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java
index e1de23393..dce451975 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java
@@ -20,6 +20,7 @@ import ai.djl.modality.Classifications;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.onnxruntime.zoo.tabular.softmax_regression.IrisFlower;
 import ai.djl.repository.zoo.Criteria;
@@ -143,6 +144,42 @@ public class OrtTest {
 
             float[][] value = (float[][]) ((OrtNDArray) ones).getTensor().getValue();
             Assert.assertEquals(value[0], new float[] {1, 1});
+
+            array = manager.create(new Shape(1), DataType.BOOLEAN);
+            Assert.assertEquals(array.getDataType(), DataType.BOOLEAN);
+
+            array = manager.create(new Shape(1), DataType.FLOAT16);
+            Assert.assertEquals(array.getDataType(), DataType.FLOAT16);
+
+            array = manager.create(new Shape(1), DataType.BFLOAT16);
+            Assert.assertEquals(array.getDataType(), DataType.BFLOAT16);
+
+            array = manager.create(new double[] {0});
+            Assert.assertEquals(array.getDataType(), DataType.FLOAT64);
+
+            array = manager.create(new Shape(1), DataType.FLOAT64);
+            Assert.assertEquals(array.getDataType(), DataType.FLOAT64);
+
+            array = manager.create(new Shape(1), DataType.INT8);
+            Assert.assertEquals(array.getDataType(), DataType.INT8);
+
+            array = manager.create(new Shape(1), DataType.UINT8);
+            Assert.assertEquals(array.getDataType(), DataType.UINT8);
+
+            array = manager.create(new int[] {0});
+            Assert.assertEquals(array.getDataType(), DataType.INT32);
+
+            array = manager.create(new Shape(1), DataType.INT32);
+            Assert.assertEquals(array.getDataType(), DataType.INT32);
+
+            array = manager.create(new long[] {0L});
+            Assert.assertEquals(array.getDataType(), DataType.INT64);
+
+            array = manager.create(new Shape(1), DataType.INT64);
+            Assert.assertEquals(array.getDataType(), DataType.INT64);
+
+            Assert.assertThrows(() -> manager.create(new Shape(0), DataType.FLOAT32));
+            Assert.assertThrows(() -> manager.create(new Shape(1), DataType.UINT32));
         }
     }