PyMC3 on Apple M1-chip

Hello! Does anybody have experience with running PyMC3 on Apple’s new M1-chip? Is it even possible?

@fonnesbeck might have some experience already?

1 Like

Looking forward to PyMC3 with a JAX backend that is optimized for the M1-chip! Is jax optimized for m1 chips on macos? · Discussion #5084 · google/jax · GitHub

I’m using PyMC3 on a daily basis on the new M1 chip. Note that Python can be run both natively and under emulation via Rosetta2; I have only had success using the latter with PyMC3, but the performance is excellent. I’ve had a few issues with multiprocessing, but I anticipate most issues are temporary as the Python scientific stack evolves to accommodate the new hardware. So, things are far from optimized for M1 yet, but I am very productive on the platform.

3 Likes

I actually already ordered a MacBook Pro with the M1 chip, so that’s good to hear :sweat_smile: Thanks for sharing your experience @fonnesbeck!

Hi, I am also considering getting a Mac with M1 chip which will be used sometimes for pymc3 computation. Any updates on the usability/compatibility so far?

Thanks!

I’m waiting for the new larger iMac (this summer?) before I switch to Apple Silicon. Glad to hear it works. Hopefully natively at some point :slight_smile:

I have pymc3 working natively. Branch here: GitHub - elbamos/pymc3: Probabilistic Programming in Python: Bayesian Modeling and Probabilistic Machine Learning with Aesara branch m1. It works with the version of aesana available on conda forge. If anyone has a test suite for benchmarking, I’m happy to run it and provide results.

3 Likes

@elbamos Could you share a little bit more how to install Aesara on Mac Silicon (M1 ARM)? My Aesara installed from miniForge 3.9 is having compilation error with simple Matrix summation example.

I didn’t do anything special - I’m running vanilla Aesara from conda forge, with Python 3.8 I believe.

I was just having trouble with this, @jjchan121, and I managed by first creating a conda environment with the appropriate requirements, and then installing pymc from the pymc3 main branch.

conda create --name pymc aesara arviz cachetools dill fastprogress pandas scipy typing-extensions
conda activate pymc
pip install https://github.com/pymc-devs/pymc3/archive/refs/heads/main.zip

One still gets a warning about missing mkl, which will impact speed, but this is a way to get it working.

New MacMini M1, I tried to install your branch (after conda create -c conda-forge -n pymc3_env python=3.9 aesara, then pip install git+https://github.com/elbamos/pymc3) but hit error that netCDF4 had conflicting dependencies:

ERROR: Cannot install pymc3 because these package versions have conflicting dependencies.

The conflict is caused by:
    arviz 0.11.4 depends on netcdf4
    arviz 0.11.3 depends on netcdf4
    arviz 0.11.2 depends on netcdf4
    arviz 0.11.1 depends on netcdf4

After conda install netcdf4, then re-running the pip install of your branch I get error:

(pymc3) ➜  ~ python
Python 3.9.7 | packaged by conda-forge | (default, Sep 29 2021, 19:22:19) 
[Clang 11.1.0 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> from netCDF4 import Dataset
>>> import pymc3 as pm
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/me/miniforge3/envs/pymc3/lib/python3.9/site-packages/pymc3/__init__.py", line 41, in <module>
    from pymc3 import gp, ode, sampling
  File "/Users/me/miniforge3/envs/pymc3/lib/python3.9/site-packages/pymc3/gp/__init__.py", line 15, in <module>
    from pymc3.gp import cov, mean, util
  File "/Users/me/miniforge3/envs/pymc3/lib/python3.9/site-packages/pymc3/gp/util.py", line 24, in <module>
    solve_lower = Solve(A_structure="lower_triangular")
TypeError: __init__() got an unexpected keyword argument 'A_structure'

My apologies if I’ve missed anything obvious or done anything daft here!

The following appears to install a functional version of the PyMC main branch without issue on M1 chipset and using Apple’s Accelerate library for BLAS:

pymc_main.yaml

name: pymc_main
channels:
  - conda-forge
dependencies:
  ## core                                                                                                                              
  - python=3.10
  - blas=*=accelerate*
  - ipykernel  ## jupyter kernel                                                                                                       

  ## Python packages                                                                                                                   
  - arviz
  - aesara
  - scipy

  ## PyPI                                                                                                                              
  - pip
  - pip:
      - -e git+ssh://git@github.com/pymc-devs/pymc.git@main#egg=pymc

creating with

mamba env create -n pymc_main -f pymc_main.yaml

Note that here the import is pymc instead of pymc3.

1 Like

@fonnesbeck: I’ve recently got a new M1 laptop from work and had a lot of issues getting my old code to work. The most common errors are: “Could not find compiler for platform METAL:”, also running parallel on pm.sample() occasionally didn’t work. I installed everything appropriately using the instruction from Apple developer forum, e.g. via apple channel and miniforge. When you did this with Rosetta2, did you run into any issues like this? Thanks :slight_smile:

Do you have the most recent XCode developer tools installed (xcode-select --install)? Do you have this problem with any PyMC model, or just yours?

You shouldn’t use Rosetta but install mambaforge which has native M1 support.

2 Likes

Hi twiecki, what’s the difference between mambaforge and miniforge since I already have the second one installed?

They work the same, you can use what you have.

I’m considering getting a mac with an M2 chip - what is the latest on ease of install / speed of pymc sampling on these chips?

I am also considering this, would be great to hear