Pretty sure it's Apple building it, and they're using JAX in-house so I imagine it will get better over time. Though they do love to drop support for old things so maybe M1 will never work again...
> Our foundation models are trained on Apple's AXLearn framework, an open-source project we released in 2023. It builds on top of JAX