[Feature request] Support bfloat16/float8 inputs in `session.run()` about onnxruntime HOT 11 OPEN

I would use a built-in array module and provide shape and type as separate argument for maximum flexibility.
However, along with that, we should have a capability to return those types as results. And right now, we only return numpy arrays.

I think OrtValue based interfaces can help us. Simpe run() returns numpy.
However, we can add methods to OrtValue to accept other types in python array and convert the result to python arrays as well.

from onnxruntime.

ORT supports bfloat16 internally. I would think a separate per type approach is not the best way to go. Rather, I would implement a new universal method that does not depend on numpy so much, and also include other types, such as quantized 8 and 4.

from onnxruntime.

I suggest using run_with_iobinding and make a helper to use run(sess, None, feeds) where is a dictionary dict[str, torch.Tensor].

from onnxruntime.

IOBinding has its own specifics as to when the input is actually copied to the device.
Also it is cumbersome for every day use.

from onnxruntime.

Thoughts on supporting dlpack protocols? This way we don't need to know if the inputs are torch or numpy or anything else. The data can be shared using the same protocol.

from onnxruntime.

Thoughts on supporting dlpack protocols? This way we don't need to know if the inputs are torch or numpy or anything else. The data can be shared using the same protocol.

Does it not come with its own set of data type enums that seem to be even more restrictive than numpy?

from onnxruntime.

Hmm. I tested it and there are constraints. dlpack supports all numpy types except for strings. It supports bfloat16 additionally, but not float8 types. Although as a ML data exchange protocol I can foresee float8 support being added.

A major advantage is that it is supported by all major frameworks (numpy, torch, jax, tensorflow, mlx etc.) so ORT can be compatible with many different tensors easily.

from onnxruntime.

import torch
import jax

t = torch.tensor(42, dtype=torch.bfloat16)
jax.dlpack.from_dlpack(t)

Array(42, dtype=bfloat16)

from onnxruntime.

We can dlpack support, in fact, I have seen some code in pybind referring to dlpack, but I am not seeing it support what we need. Should we not design something that we can own, maintain and deliver what we need?

from onnxruntime.

It seems, we can extend numpy types support for our needs. See for example a project that implements bfloat16 type and makes it 1st class citizen for numpy. https://github.com/GreenWaves-Technologies/bfloat16/tree/main

from onnxruntime.

There is also https://github.com/jax-ml/ml_dtypes (from Google) that implements all relevant types, which we use in the ONNX IR. Upstreaming seems challenging based on the context this project gave.

from onnxruntime.