Initial support

.pre-commit-config.yaml

@@ -0,0 +1,11 @@
+repos:
+-   repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 24.2.0
+    hooks:
+    -   id: black
+-   repo: https://github.com/pycqa/isort
+    rev: 5.13.2
+    hooks:
+    -   id: isort
+        args:
+            - --profile=black

README.md

@@ -1,2 +1,25 @@
-# mlx-onnx
-MLX support for the Open Neural Network Exchange (ONNX)
+# MLX ONNX
+
+MLX support for the Open Neural Network Exchange ([ONNX](https://onnx.ai/)) 
+
+## Install
+
+```shell
+pip install mlx-onnx
+```
+
+## Usage
+
+```python
+from mlx.onnx import MlxBackend
+from onnx import hub
+
+model = hub.load("mnist")
+backend = MlxBackend(model)
+result = backend.run(...) # pass inputs to model
+```
+
+## Examples
+
+- [ResNet](./examples/resnet/example.py)
+- [Mnist](./examples/mnist/example.py)

examples/mnist/example.py

@@ -0,0 +1,19 @@
+# Copyright © 2024 Apple Inc.
+
+import mlx.core as mx
+import numpy as np
+from onnx import hub
+from PIL import Image
+
+from mlx.onnx import MlxBackend
+
+if __name__ == "__main__":
+    x = (
+        mx.array(np.asarray(Image.open("./nine.jpeg")))
+        .reshape((1, 1, 28, 28))
+        .astype(mx.float32)
+    )
+    model = hub.load("mnist")
+    backend = MlxBackend(model)
+    res = backend.run(x)
+    print(f"It was a {mx.argmax(res[0]).item()}")

examples/resnet/example.py

@@ -0,0 +1,33 @@
+# Copyright © 2024 Apple Inc.
+
+import mlx.core as mx
+import mlx.data as dx
+import onnx
+
+from mlx.onnx import MlxBackend
+
+
+def run(image: str):
+    dataset = (
+        dx.buffer_from_vector([{"file_name": image.encode()}])
+        .load_image("file_name", output_key="image")
+        .image_resize_smallest_side("image", 256)
+        .image_center_crop("image", 224, 224)
+        .key_transform("image", lambda x: (x - 127.0) / 128.0)
+    )
+    with open("./imagenet_labels.txt") as f:
+        labels = [l.strip() for l in f.readlines()]
+
+    model = onnx.hub.load("resnet50")
+    backend = MlxBackend(model)
+    res = []
+    for data in dataset:
+        img = mx.array(data["image"]).transpose(2, 0, 1)[None]
+        x = backend.run(img)[0]
+        res.append((labels[mx.argmax(x).item()], mx.max(x).item()))
+    return res
+
+
+if __name__ == "__main__":
+    for label, score in run("./car.jpg"):
+        print(f"Image containes a {label} with score {score:.3f}.")