mirror of
https://github.com/priyanshujain/torchmlx.git
synced 2026-10-02 11:07:13 +00:00
563 B
563 B
torchmlx
torchmlx is a pytorch-shaped compatibility layer that uses mlx on apple silicon and pytorch elsewhere. it is experimental and built first for educational use.
pip install pytorchmlx
import torchmlx as torch
from torchmlx import nn, optim
see the tinystories example and compatibility details.
| backend | 5,000 steps + generation | effective steps/s |
|---|---|---|
| torchmlx mlx, m3 pro | 170.27 s | 29.4 |
| pytorch mps, m3 pro | 184.42 s | 27.1 |