saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[java] Allow the creation of boolean tensors from ByteBuffer (#15556)

Browse source 
### 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.
This commit is contained in:
Adam Pocock 2023-06-05 12:58:50 -04:00 committed by GitHub
parent a95f8ae53c
commit 3c2a11f2f1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 45 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
java/src/main/java/ai/onnxruntime/OrtUtil.java
View file

@ -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();

40
java/src/test/java/ai/onnxruntime/InferenceTest.java
View file

@ -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();
}
}