Apple supports JAX[0] along with PyTorch[1] and Tensorflow[2] on macOS with both... | Hacker News

Hacker Newsnew | past | comments | ask | show | jobs | submit

		woadwarrior01 on Sept 29, 2023 \| parent \| context \| favorite \| on: JAX – NumPy on the CPU, GPU, and TPU Apple supports JAX[0] along with PyTorch[1] and Tensorflow[2] on macOS with both Apple Silicon and AMD GPUs (on x86 Macs). Although, the perf isn't great. I write most of my experimental ML code in JAX on an M2 Macbook Air and then move to a proper multi-GPU Linux box for full training runs. [0]: https://developer.apple.com/metal/jax/ [1]: https://developer.apple.com/metal/pytorch/ [2]: https://developer.apple.com/metal/tensorflow-plugin/

VHRanger on Sept 29, 2023 [–]

Pytorch on my M2 max using the MPS backend has pretty decent performance to be honest?

It's significantly faster than CPU. Something like 100x using sheet

Guidelines | FAQ | Lists | API | Security | Legal | Apply to YC | Contact