mlx icon indicating copy to clipboard operation
mlx copied to clipboard

[Feature Request] ONNX runtime support

Open dc-dc-dc opened this issue 2 years ago • 3 comments

ONNX is an open-standard for machine learning interoperability. Allowing developers to utilize different backends to execute their models.

The backend implementation is relatively straight forward where we would need to provide mappings from their ops to the mlx ops. They also provide a testing suite that makes it simple to exclude any unsupported ops. Should support be added directly to mlx similar to how pytorch has their torch.onnx module or a separate repo like jaxonnxruntime

dc-dc-dc avatar Dec 26 '23 16:12 dc-dc-dc

Dup of #16 (I'll close 16 in favor of discussing here).

@dc-dc-dc yes I think onnx support would be great! Regarding support directly in mlx vs a separate repo / runtime I am not sure. Any idea why Jax went with the latter?

awni avatar Dec 27 '23 14:12 awni

Did some more digging into this, looks like pytorch only supports exporting to onnx even though importing/running was a requested feature. I don't have anything concrete on the decision from jax to make it separate but tensorflow seems to have gone this route as well having onnx-tf . Though there are some libs like tinygrad which added support for the runtime directly.

dc-dc-dc avatar Dec 27 '23 20:12 dc-dc-dc

Also would the implementation be in the core C++ lib or in python, the onnx python package provides the backend helper to create the runtime. If C++ implementation is desired would need a protobuf parser i believe.

dc-dc-dc avatar Dec 28 '23 16:12 dc-dc-dc