search

apple / tensorflow_macos

TensorFlow for macOS 11.0+ accelerated using Apple's ML Compute framework.
Other
3.66k stars 308 forks source link

JAX support ML Compute #180

Open tetterl opened 3 years ago

tetterl commented 3 years ago

Are there any plans to also support JAX through the ML Compute framework? Since JAX gets compiled to XLA which should already work there shouldn't be too many things to adjust?

See: https://github.com/google/jax/issues/5501

machineko commented 3 years ago

+1 for that excellent idea :)

DavidePaglieri commented 3 years ago

It would be very helpful to know whether there are any plans for this. My decision on buying an M1 mac or sticking to windows kind of depends on this at the moment.