Installation: apple M2

[ChauhanT]

Hi,

Thank you very much for sharing your wonderful work with the community.

I am trying to install brainpy on my apple M2 machine (virtual environment python 3.10.13). I can install everything fine, but when I import brainpy, I get this error:

ImportError: cannot import name 'Primitive' from 'jax.core' (/Users/USER/Work/dev/ei_srp/.venv/lib/python3.10/site-packages/jax/core.py)

I thought installing jax-metal would do the trick, but it didn't. My pyproject.toml dependencies read as follows:

requires-python = ">=3.10,<3.11"
dependencies = [
    "brainpy>=2.6.0.post20241205",
    "jax-metal>=0.1.1",
]

Any tips/advice would be awesome as I can't seem to figure it out by forum-browsing and googling.

Cheers!

[Routhleck]

Hi, Thanks for your feedback.
Do not install jax-metal. This package was integrated into JAX a long time ago—now you can simply install jax directly to use Apple Silicon (M1/M2) chips. I believe the issue is caused by using a too-recent version of JAX. Please try installing an older version, such as:

pip install jax==0.4.35

[ChauhanT]

That does indeed fix the install. Thank you for the prompt response.

To help others who may experience this, here is the relevant part of my pyproject.toml :

requires-python = ">=3.10,<3.11"
dependencies = [
    "brainpy>=2.6.0.post20241205",
    "jax==0.4.35",
]