AstroNN is a lovely Python package for applying deep learning to astronomy: stellar spectra, Gaia data, the works. Getting it and TensorFlow running on Apple Silicon took some fiddling; here's the setup that finally worked.

The problem

Stock TensorFlow wheels don't target the M-series GPU, so a naive pip install tensorflow either fails to build or runs CPU-only and slow. The fix is Apple's tensorflow-metal plugin on top of tensorflow-macos.

It's worth knowing what that plugin actually is, because it explains the failure modes. TensorFlow exposes a PluggableDevice interface, which lets a vendor register a new accelerator backend without TensorFlow itself knowing anything about the hardware. Apple's plugin registers a device it calls GPU and routes the kernels it implements to Metal Performance Shaders. Anything it does not implement silently falls back to the CPU. Nothing errors. Your model just runs at a tenth of the speed you expected and you have no idea why.

The other quiet constraint is precision. The Metal path is a float32 world. If you set a float64 global policy, which is tempting for anything astronomical where you are used to caring about precision, the ops drop back to CPU and the GPU sits idle.

The recipe

conda create -n astronn python=3.10
conda activate astronn
pip install tensorflow-macos tensorflow-metal
pip install astroNN

Keep Python at 3.10. Newer versions raced ahead of the metal plugin's support matrix when I set this up. Pin the pair of TensorFlow packages to matching minor versions too; tensorflow-metal is built against a specific tensorflow-macos, and a mismatch surfaces as an import-time symbol error rather than anything helpful.

After that, tf.config.list_physical_devices('GPU') shows the device and models train on the GPU. That check only proves the plugin loaded, though. To prove work is actually landing there:

tf.debugging.set_log_device_placement(True)

which prints the device chosen for each op, and will show you exactly which layer is quietly executing on the CPU.

Worth it?

For experimenting with pretrained astro models on a laptop, absolutely: no cloud GPU bill, and the M-series chips are no slouch. The unified memory helps more than the raw throughput does, since the whole of system RAM is addressable by the GPU and you are not juggling a 12 GB card. Large batches of spectra fit where they wouldn't on a mid-range discrete GPU.