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.