search

ml-explore / mlx

MLX: An array framework for Apple silicon
https://ml-explore.github.io/mlx/
MIT License
16.37k stars 931 forks source link

[Feature Request] ONNX runtime support #290

Open dc-dc-dc opened 8 months ago

dc-dc-dc commented 8 months ago

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

awni commented 8 months ago

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?

dc-dc-dc commented 8 months ago

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 commented 8 months ago

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.