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.

2 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