macOS arm64 build crashes with "Symbol not found"

[dfm]

• [x] I read the conda-forge documentation and could not find the solution for my problem there.

Issue:

Thanks for your work with the macOS arm64 build @xhochy - I'm really looking forward to getting this to work!

On my new mac, when I create a fresh environment:

conda create -n jaxlib-test -c conda-forge python=3.9 jax jaxlib

and then import jax (or jaxlib.xla_client), it fails with the following ImportError:

Python 3.9.6 | packaged by conda-forge | (default, Jul 11 2021, 03:35:11) 
[Clang 11.1.0 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
/opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test/lib/python3.9/site-packages/jax/lib/__init__.py:31: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.
  warnings.warn("JAX on Mac ARM machines is experimental and minimally tested. "
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test/lib/python3.9/site-packages/jax/__init__.py", line 37, in <module>
    from . import config as _config_module
  File "/opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test/lib/python3.9/site-packages/jax/config.py", line 18, in <module>
    from jax._src.config import config
  File "/opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test/lib/python3.9/site-packages/jax/_src/config.py", line 27, in <module>
    from jax import lib
  File "/opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test/lib/python3.9/site-packages/jax/lib/__init__.py", line 73, in <module>
    from jaxlib import xla_client
  File "/opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test/lib/python3.9/site-packages/jaxlib/xla_client.py", line 31, in <module>
    from . import xla_extension as _xla
ImportError: dlopen(/opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test/lib/python3.9/site-packages/jaxlib/xla_extension.so, 2): Symbol not found: __ZN3jax12_GLOBAL__N_121CompiledFunctionCache16kDefaultCapacityE
  Referenced from: /opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test/lib/python3.9/site-packages/jaxlib/xla_extension.so
  Expected in: flat namespace
 in /opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test/lib/python3.9/site-packages/jaxlib/xla_extension.so

Do y'all have any thoughts or suggestions?

Environment (conda list):

``` $ conda list # packages in environment at /opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test: # # Name Version Build Channel absl-py 0.13.0 pyhd8ed1ab_0 conda-forge ca-certificates 2021.5.30 h4653dfc_0 conda-forge jax 0.2.19 pyhd8ed1ab_0 conda-forge jaxlib 0.1.69 py39h7f752ed_0 conda-forge libblas 3.9.0 11_osxarm64_openblas conda-forge libcblas 3.9.0 11_osxarm64_openblas conda-forge libcxx 12.0.1 h168391b_0 conda-forge libgfortran 5.0.0.dev0 11_0_1_hf114ba7_23 conda-forge libgfortran5 11.0.1.dev0 hf114ba7_23 conda-forge liblapack 3.9.0 11_osxarm64_openblas conda-forge libopenblas 0.3.17 openmp_h5dd58f0_1 conda-forge llvm-openmp 12.0.1 hf3c4609_1 conda-forge ncurses 6.2 h9aa5885_4 conda-forge numpy 1.21.2 py39h1f3b974_0 conda-forge openssl 1.1.1k h3422bc3_1 conda-forge opt_einsum 3.3.0 pyhd8ed1ab_1 conda-forge pip 21.2.4 pyhd8ed1ab_0 conda-forge python 3.9.6 h54d631c_1_cpython conda-forge python-flatbuffers 2.0 pyhd8ed1ab_0 conda-forge python_abi 3.9 2_cp39 conda-forge readline 8.1 hedafd6a_0 conda-forge scipy 1.7.0 py39h5060c3b_0 conda-forge setuptools 57.4.0 py39h2804cbe_0 conda-forge six 1.16.0 pyh6c4a22f_0 conda-forge sqlite 3.36.0 h72a2b83_0 conda-forge tk 8.6.11 he1e0b03_0 conda-forge tzdata 2021a he74cb21_1 conda-forge wheel 0.37.0 pyhd8ed1ab_1 conda-forge xz 5.2.5 h642e427_1 conda-forge zlib 1.2.11 h31e879b_1009 conda-forge ```

Details about conda and system ( conda info ):

``` $ conda info active environment : jaxlib-test active env location : /opt/homebrew/Caskroom/miniforge/base/envs/jaxlib-test shell level : 3 user config file : /Users/dforemanmackey/.condarc populated config files : /opt/homebrew/Caskroom/miniforge/base/.condarc conda version : 4.10.3 conda-build version : not installed python version : 3.9.6.final.0 virtual packages : __osx=11.5.2=0 __unix=0=0 __archspec=1=arm64 base environment : /opt/homebrew/Caskroom/miniforge/base (writable) conda av data dir : /opt/homebrew/Caskroom/miniforge/base/etc/conda conda av metadata url : None channel URLs : https://conda.anaconda.org/conda-forge/osx-arm64 https://conda.anaconda.org/conda-forge/noarch package cache : /opt/homebrew/Caskroom/miniforge/base/pkgs /Users/dforemanmackey/.conda/pkgs envs directories : /opt/homebrew/Caskroom/miniforge/base/envs /Users/dforemanmackey/.conda/envs platform : osx-arm64 user-agent : conda/4.10.3 requests/2.26.0 CPython/3.9.6 Darwin/20.6.0 OSX/11.5.2 UID:GID : 1185:20 netrc file : None offline mode : False ```

[xhochy]

Verified locally that #64 fixes this.

[xhochy]

The master builds still has this issue while it wasn't present in the PR.

[xhochy]

This issue here is that the xla code requires C++17 and we are building with -std=c++14.

[beckermr]

I am seeing this on the x85 osx builds too. #70

[xhochy]

x86 should be solved by https://github.com/conda-forge/jaxlib-feedstock/pull/65, I was just reluctant to merge that as it doesn't solve all of the osx-arm64 problems (but does fix them all for osx-arm64)

[beckermr]

I am very appreciative of the work here and so merge at your own pace. I am using 0.1.69 for now.

[beckermr]

We should make a pr to mark versions at it above 0.70 broken for osx too.