From 486b07bb07af7508d01cd0b45435f0358ddd3d2f Mon Sep 17 00:00:00 2001 From: Toma Susi Date: Thu, 1 Oct 2026 10:48:36 +0200 Subject: [PATCH] Drop the enable_mps step from the Metal installation instructions abTEM/abTEM#414 now imports PyTorch on the first use of the 'mps' device instead of behind an 'enable_mps' flag set before abTEM is imported, so the flag, its environment variable and the note explaining why abtem.config.set could not turn it on all go. In their place: making 'mps' the default device, and the one ordering caveat left -- abTEM (or torch) has to be imported before pyfftw. Co-Authored-By: Claude Opus 5.5 --- docs/getting_started/install.md | 31 ++++++++++++++----------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/docs/getting_started/install.md b/docs/getting_started/install.md index c9dbc930..a92b86e6 100644 --- a/docs/getting_started/install.md +++ b/docs/getting_started/install.md @@ -135,23 +135,8 @@ pip install "abtem[mps]" The extra is restricted to macOS on Apple silicon, since Metal exists nowhere else. -To enable the backend, set `enable_mps` **before** *ab*TEM is imported — either through the environment: - -```{code-block} -ABTEM_ENABLE_MPS=true python your_script.py -``` - -or in `~/.config/abtem/abtem.yaml`: - -```{code-block} -enable_mps: true -``` - -```{note} -`abtem.config.set(enable_mps=True)` does **not** work: the setting decides whether PyTorch is imported before -FFTW, and by the time the call runs that has already been settled. PyTorch and FFTW each bundle their own copy -of `libomp`, and a process that loaded FFTW's first crashes inside ordinary PyTorch operations. -``` +Nothing needs enabling: *ab*TEM imports PyTorch the first time the `mps` device is used, so importing *ab*TEM +stays as fast as without it. You can verify that Metal support is available using the code below: @@ -171,6 +156,18 @@ potential = abtem.Potential(atoms, gpts=512, device="mps") probe = abtem.Probe(energy=200e3, semiangle_cutoff=20, device="mps") ``` +or make Metal the default device, with `abtem.config.set({"device": "mps"})` or in `~/.config/abtem/abtem.yaml`: + +```{code-block} +device: mps +``` + +```{note} +Import *ab*TEM (or PyTorch) before importing `pyfftw` yourself. PyTorch and FFTW each bundle their own copy of +`libomp`, and PyTorch's has to be loaded first; *ab*TEM arranges that when it is imported. If `pyfftw` was +imported earlier, using the `mps` device raises an error explaining this, rather than letting PyTorch crash. +``` + Two things are worth knowing before you benchmark: - **Metal is single precision.** `precision` must be `float32`; pairing `device="mps"` with `float64` is