From 3c2a11f2f1d8aeaf30f47402a5598dbd30990e9c Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Mon, 5 Jun 2023 12:58:50 -0400
Subject: [PATCH] [java] Allow the creation of boolean tensors from ByteBuffer
 (#15556)

### Description
The tensor creation code now allows the creation of boolean tensors from
non-direct `ByteBuffer` instances. It previously only allowed them from
arrays and direct `ByteBuffer` instances and this fixes that
inconsistency. The boolean tensor test has been updated to cover all
three cases.

### Motivation and Context
Fixes #15509.
---
 .../src/main/java/ai/onnxruntime/OrtUtil.java | 10 +++--
 .../java/ai/onnxruntime/InferenceTest.java    | 40 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/java/src/main/java/ai/onnxruntime/OrtUtil.java b/java/src/main/java/ai/onnxruntime/OrtUtil.java
index eb27d1dafd..e6f5bc1f5d 100644
--- a/java/src/main/java/ai/onnxruntime/OrtUtil.java
+++ b/java/src/main/java/ai/onnxruntime/OrtUtil.java
@@ -493,6 +493,9 @@ public final class OrtUtil {
    * @return The prepared buffer tuple.
    */
   static BufferTuple prepareBuffer(Buffer data, OnnxJavaType type) {
+    if (type == OnnxJavaType.STRING || type == OnnxJavaType.UNKNOWN) {
+      throw new IllegalStateException("Cannot create a " + type + " tensor from a buffer");
+    }
     int bufferPos;
     long bufferSizeLong = data.remaining() * (long) type.size;
     if (bufferSizeLong > (Integer.MAX_VALUE - (8 * type.size))) {
@@ -522,6 +525,7 @@ public final class OrtUtil {
         case DOUBLE:
           tmp = buffer.asDoubleBuffer().put((DoubleBuffer) data);
           break;
+        case BOOL:
         case UINT8:
         case INT8:
           // buffer is already a ByteBuffer, no cast needed.
@@ -536,12 +540,10 @@ public final class OrtUtil {
         case INT64:
           tmp = buffer.asLongBuffer().put((LongBuffer) data);
           break;
-        case BOOL:
-        case STRING:
-        case UNKNOWN:
         default:
           throw new IllegalStateException(
-              "Impossible to reach here, managed to cast a buffer as an incorrect type");
+              "Impossible to reach here, managed to cast a buffer as an incorrect type, found "
+                  + type);
       }
       data.position(origPosition);
       tmp.rewind();
diff --git a/java/src/test/java/ai/onnxruntime/InferenceTest.java b/java/src/test/java/ai/onnxruntime/InferenceTest.java
index 8fc5fe2daf..36ae3cc356 100644
--- a/java/src/test/java/ai/onnxruntime/InferenceTest.java
+++ b/java/src/test/java/ai/onnxruntime/InferenceTest.java
@@ -1168,8 +1168,11 @@ public class InferenceTest {
         OrtSession session = env.createSession(modelPath, options)) {
       String inputName = session.getInputNames().iterator().next();
       Map<String, OnnxTensor> container = new HashMap<>();
+      long[] shape = new long[] {1, 5};
+
+      // Test array input
       boolean[] flatInput = new boolean[] {true, false, true, false, true};
-      Object tensorIn = OrtUtil.reshape(flatInput, new long[] {1, 5});
+      Object tensorIn = OrtUtil.reshape(flatInput, shape);
       OnnxTensor ov = OnnxTensor.createTensor(env, tensorIn);
       container.put(inputName, ov);
       try (OrtSession.Result res = session.run(container)) {
@@ -1177,6 +1180,41 @@ public class InferenceTest {
         assertArrayEquals(flatInput, resultArray);
       }
       OnnxValue.close(container);
+      container.clear();
+
+      // Test direct buffer input
+      ByteBuffer dirBuf = ByteBuffer.allocateDirect(5).order(ByteOrder.nativeOrder());
+      dirBuf.put((byte) 1);
+      dirBuf.put((byte) 0);
+      dirBuf.put((byte) 1);
+      dirBuf.put((byte) 0);
+      dirBuf.put((byte) 1);
+      dirBuf.rewind();
+      ov = OnnxTensor.createTensor(env, dirBuf, shape, OnnxJavaType.BOOL);
+      container.put(inputName, ov);
+      try (OrtSession.Result res = session.run(container)) {
+        boolean[] resultArray = TestHelpers.flattenBoolean(res.get(0).getValue());
+        assertArrayEquals(flatInput, resultArray);
+      }
+      OnnxValue.close(container);
+      container.clear();
+
+      // Test non-direct buffer input
+      ByteBuffer buf = ByteBuffer.allocate(5);
+      buf.put((byte) 1);
+      buf.put((byte) 0);
+      buf.put((byte) 1);
+      buf.put((byte) 0);
+      buf.put((byte) 1);
+      buf.rewind();
+      ov = OnnxTensor.createTensor(env, buf, shape, OnnxJavaType.BOOL);
+      container.put(inputName, ov);
+      try (OrtSession.Result res = session.run(container)) {
+        boolean[] resultArray = TestHelpers.flattenBoolean(res.get(0).getValue());
+        assertArrayEquals(flatInput, resultArray);
+      }
+      OnnxValue.close(container);
+      container.clear();
     }
   }