如何在M1 Mac上安装macOS 12上的trax、jax、jaxlib？

Question

trax

对于trax，我尝试在本地运行它(macOS 12.1，苹果硅ARM M1处理器，8GB内存，Anaconda)，但我遇到了一些问题。

在python3.8.5环境中，我在(Anaconda) conda环境中安装了运行trax的pip3 install trax==1.3.9。后来，当我尝试用trax导入代码中的from trax import layers as tl层时，遇到了一些问题。

RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. You may be able work around this issue by building jaxlib from source.

我匆忙地用python 3.10启动了一个新的conda环境。然后在试图安装trax时遇到了问题

macos ERROR: Could not find a version that satisfies the requirement tensorflow-text (from trax) (from versions: none)
ERROR: No matching distribution found for tensorflow-text

然后，我用python3.9创建了一个新环境。安装进行得很好，但是导入层时遇到了相同的错误：

RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. You may be able work around this issue by building jaxlib from source.

在jaxlib和jax之前安装trax

然后，我尝试在这些指示之后的源代码中构建jaxlib，并得到了以下错误：

Building XLA and installing it in the jaxlib source tree...
./bazel-4.2.1-darwin-x86_64 run --verbose_failures=true --config=avx_posix --config=mkl_open_source_only :build_wheel -- --output_path=/my path/jax/dist --cpu=x86_64
ERROR: bazel does not currently work properly from paths containing spaces (/my path/jax).
b''
Traceback (most recent call last):
  File "/my path/jax/build/build.py", line 524, in <module>
    main()
  File "/my path/jax/build/build.py", line 519, in main
    shell(command)
  File "/my path/jax/build/build.py", line 53, in shell
    output = subprocess.check_output(cmd)
  File "/myuserpath/opt/anaconda3/envs/mytraxenv/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/myuserpath/opt/anaconda3/envs/mytraxenv/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-4.2.1-darwin-x86_64', 'run', '--verbose_failures=true', '--config=avx_posix', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/my path/jax/dist', '--cpu=x86_64']' returned non-zero exit status 36.

强调我最初忽略的部分: bazel目前在包含空格的路径(/my path/jax)中不能正常工作。

我将/my path/目录移动到没有空格/mypath/的路径上。删除并重新下载jax目录。尽管如此，使用python build/build.py构建(用于CPU)还是失败了：

ERROR: /private/var/tmp/_bazel_a/2caf512c3c5e3f3f654bc58b48b8333a/external/llvm-project/llvm/BUILD.bazel:610:11: Generating code from table: include/llvm/IR/Intrinsics.td @llvm-project//llvm:intrinsic_XCore_gen__gen_intrinsic_enums__intrinsic_prefix_xcore_genrule failed: (Illegal instruction): bash failed: error executing command 
  (cd /private/var/tmp/_bazel_a/2caf512c3c5e3f3f654bc58b48b8333a/execroot/__main__ && \
  exec env - \
    PATH=/myuserpath/opt/anaconda3/envs/mytraxenv/bin:/myuserpath/opt/anaconda3/condabin:/opt/homebrew/bin:/opt/homebrew/sbin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin:/Library/TeX/texbin:/Library/Apple/usr/bin \
  /bin/bash -c 'source external/bazel_tools/tools/genrule/genrule-setup.sh; bazel-out/darwin-opt/bin/external/llvm-project/llvm/llvm-tblgen -I external/llvm-project/llvm/include -I external/llvm-project/clang/include -I $(dirname external/llvm-project/llvm/include/llvm/IR/Intrinsics.td) -gen-intrinsic-enums -intrinsic-prefix=xcore external/llvm-project/llvm/include/llvm/IR/Intrinsics.td  -o bazel-out/darwin-opt/bin/external/llvm-project/llvm/include/llvm/IR/IntrinsicsXCore.h')
Execution platform: @local_execution_config_platform//:platform
/bin/bash: line 1: 11140 Illegal instruction: 4  bazel-out/darwin-opt/bin/external/llvm-project/llvm/llvm-tblgen -I external/llvm-project/llvm/include -I external/llvm-project/clang/include -I $(dirname external/llvm-project/llvm/include/llvm/IR/Intrinsics.td) -gen-intrinsic-enums -intrinsic-prefix=xcore external/llvm-project/llvm/include/llvm/IR/Intrinsics.td -o bazel-out/darwin-opt/bin/external/llvm-project/llvm/include/llvm/IR/IntrinsicsXCore.h
Target //build:build_wheel failed to build
INFO: Elapsed time: 620.950s, Critical Path: 45.35s
INFO: 589 processes: 132 internal, 457 local.
FAILED: Build did NOT complete successfully
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully
b''
Traceback (most recent call last):
  File "/mypath/jax/build/build.py", line 524, in <module>
    main()
  File "/mypath/jax/build/build.py", line 519, in main
    shell(command)
  File "/mypath/jax/build/build.py", line 53, in shell
    output = subprocess.check_output(cmd)
  File "/myuserpath/opt/anaconda3/envs/mytraxenv/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/myuserpath/opt/anaconda3/envs/mytraxenv/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-4.2.1-darwin-x86_64', 'run', '--verbose_failures=true', '--config=avx_posix', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/mypath/jax/dist', '--cpu=x86_64']' returned non-zero exit status 1.

再试几次(删除并重新下载了jax目录)和.../lib/python3.9/subprocess.py中的第528行似乎导致了问题，但是输出(虽然大部分是上面的输出)有时略有不同。这让我对记忆产生了怀疑，因为我(诚然)已经几个星期没有重新启动我的机器了，而且它开始毫无反应地慢下来。

我将我的XCode命令行工具更新为12.2版(非常肯定)。

我重新启动了我的(8GB)机器。删除并重新下载jax目录。我安装了bazel版本5.0.0与自制，以防止这将有所帮助。我有点担心它一直在为我的ARM处理器下载一个x86版本。bazel安装进行得很顺利。

再次从这些指示开始。然而，jaxlib构建清楚地表明，它想要一个更早的版本(4.2.1) bazel，并像以前一样下载了它：

b'\x1b[31mERROR: The project you\'re trying to build requires Bazel 4.2.1 (specified in /mypath/jax/.bazelversion), but it wasn\'t found in /opt/homebrew/Cellar/bazel/5.0.0/libexec/bin.\x1b[0m\n\nBazel binaries for all official releases can be downloaded from here:\n  https://github.com/bazelbuild/bazel/releases\n\nYou can download the required version directly using this command:\n  (cd "/opt/homebrew/Cellar/bazel/5.0.0/libexec/bin" && curl -fLO https://releases.bazel.build/4.2.1/release/bazel-4.2.1-darwin-x86_64 && chmod +x bazel-4.2.1-darwin-x86_64)\n'

再次，在同一行528中出现了不同的错误。现在显示全速运行：

b'\x1b[31mERROR: The project you\'re trying to build requires Bazel 4.2.1 (specified in /mypath/jax/.bazelversion), but it wasn\'t found in /opt/homebrew/Cellar/bazel/5.0.0/libexec/bin.\x1b[0m\n\nBazel binaries for all official releases can be downloaded from here:\n  https://github.com/bazelbuild/bazel/releases\n\nYou can download the required version directly using this command:\n  (cd "/opt/homebrew/Cellar/bazel/5.0.0/libexec/bin" && curl -fLO https://releases.bazel.build/4.2.1/release/bazel-4.2.1-darwin-x86_64 && chmod +x bazel-4.2.1-darwin-x86_64)\n'
Downloading bazel from: https://github.com/bazelbuild/bazel/releases/download/4.2.1/bazel-4.2.1-darwin-x86_64
bazel-4.2.1-darwin-x86_64 [########################################] 100%
Bazel binary path: ./bazel-4.2.1-darwin-x86_64
Bazel version: 4.2.1
Python binary path: /myuserpath/opt/anaconda3/envs/mytraxenv/bin/python
Python version: 3.9
NumPy version: 1.21.2
MKL-DNN enabled: yes
Target CPU: x86_64
Target CPU features: release
CUDA enabled: no
TPU enabled: no
ROCm enabled: no

Building XLA and installing it in the jaxlib source tree...
./bazel-4.2.1-darwin-x86_64 run --verbose_failures=true --config=avx_posix --config=mkl_open_source_only :build_wheel -- --output_path=/mypath/jax/dist --cpu=x86_64
INFO: Options provided by the client:
  Inherited 'common' options: --isatty=0 --terminal_columns=80
INFO: Reading rc options for 'run' from /mypath/jax/.bazelrc:
  Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /mypath/jax/.bazelrc:
  Inherited 'build' options: --apple_platform_type=macos --macos_minimum_os=10.9 --announce_rc --define open_source_build=true --spawn_strategy=standalone --enable_platform_specific_config --define=no_aws_support=true --define=no_gcp_support=true --define=no_hdfs_support=true --define=no_kafka_support=true --define=no_ignite_support=true --define=grpc_no_ares=true -c opt --config=short_logs --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.
INFO: Reading rc options for 'run' from /mypath/jax/.jax_configure.bazelrc:
  Inherited 'build' options: --strategy=Genrule=standalone --repo_env PYTHON_BIN_PATH=/myuserpath/opt/anaconda3/envs/mytraxenv/bin/python --action_env=PYENV_ROOT --python_path=/myuserpath/opt/anaconda3/envs/mytraxenv/bin/python --distinct_host_configuration=false
INFO: Found applicable config definition build:short_logs in file /mypath/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:avx_posix in file /mypath/jax/.bazelrc: --copt=-mavx --host_copt=-mavx
INFO: Found applicable config definition build:mkl_open_source_only in file /mypath/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:macos in file /mypath/jax/.bazelrc: --config=posix
INFO: Found applicable config definition build:posix in file /mypath/jax/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
Loading: 
Loading: 0 packages loaded
Analyzing: target //build:build_wheel (1 packages loaded, 0 targets configured)
Analyzing: target //build:build_wheel (8 packages loaded, 286 targets configured)
Analyzing: target //build:build_wheel (10 packages loaded, 4271 targets configured)
Analyzing: target //build:build_wheel (10 packages loaded, 4271 targets configured)
Analyzing: target //build:build_wheel (11 packages loaded, 4494 targets configured)
INFO: Analyzed target //build:build_wheel (11 packages loaded, 7417 targets configured).

INFO: Found 1 target...
[0 / 3] [Prepa] BazelWorkspaceStatusAction stable-status.txt
[21 / 233] Compiling src/google/protobuf/generated_enum_util.cc; 2s local ... (6 actions, 3 running)
[28 / 233] Compiling src/google/protobuf/extension_set.cc; 4s local ... (8 actions, 7 running)
[114 / 523] Compiling src/google/protobuf/generated_message_util.cc; 2s local ... (8 actions running)
[140 / 594] Compiling platform/c++11/src/nsync_semaphore_mutex.cc; 0s local ... (8 actions, 7 running)
[237 / 597] Compiling src/google/protobuf/util/message_differencer.cc; 0s local ... (8 actions running)
[252 / 597] Compiling src/google/protobuf/util/message_differencer.cc; 4s local ... (8 actions running)
[272 / 597] Compiling src/google/protobuf/descriptor.pb.cc; 8s local ... (8 actions running)
[296 / 597] Compiling src/google/protobuf/descriptor.cc; 12s local ... (8 actions running)
[315 / 597] Compiling src/google/protobuf/descriptor_database.cc; 2s local ... (5 actions running)
[487 / 1,963] Compiling src/compiler/python_generator.cc; 2s local ... (8 actions running)
[579 / 3,117] Compiling tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv3d.cc; 0s local ... (8 actions, 7 running)
[619 / 3,173] Compiling tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv3d.cc; 8s local ... (8 actions running)
[687 / 3,270] Compiling tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv3d.cc; 17s local ... (8 actions running)
[774 / 3,464] Compiling tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.cc; 27s local ... (8 actions running)
[1,204 / 4,860] Compiling tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.cc; 39s local ... (8 actions running)
[1,255 / 4,916] Compiling tensorflow/compiler/xla/service/cpu/runtime_matmul.cc; 52s local ... (8 actions running)
[1,340 / 5,042] Compiling tensorflow/compiler/xla/service/cpu/runtime_matmul.cc; 68s local ... (8 actions running)
[1,456 / 5,156] Compiling tensorflow/compiler/xla/service/cpu/runtime_matmul.cc; 86s local ... (8 actions, 7 running)
[1,661 / 5,704] Compiling tensorflow/compiler/xla/service/cpu/runtime_matmul.cc; 107s local ... (8 actions, 7 running)
[1,688 / 5,704] Compiling tensorflow/compiler/xla/service/cpu/runtime_matmul.cc; 132s local ... (8 actions, 7 running)
[1,721 / 5,704] Compiling tensorflow/compiler/xla/service/cpu/runtime_matmul.cc; 160s local ... (8 actions, 7 running)
[2,106 / 6,584] Compiling tensorflow/compiler/xla/service/cpu/runtime_matmul.cc; 192s local ... (8 actions, 7 running)
[2,342 / 7,067] Compiling tensorflow/compiler/xla/service/cpu/runtime_matmul.cc; 231s local ... (8 actions, 7 running)
[2,378 / 7,067] Compiling tensorflow/compiler/xla/service/cpu/runtime_matmul.cc; 274s local ... (8 actions, 7 running)
[2,422 / 7,067] Compiling src/cpu/rnn/ref_rnn.cpp; 75s local ... (8 actions, 7 running)
[2,452 / 7,067] Compiling src/cpu/x64/gemm/f32/jit_avx512_core_f32_copy_at_kern_part1_autogen.cpp; 54s local ... (8 actions, 7 running)
[2,500 / 7,067] Compiling src/cpu/x64/gemm/f32/jit_avx512_core_f32_copy_at_kern_part1_autogen.cpp; 119s local ... (8 actions, 7 running)
[2,577 / 7,067] Compiling src/common/memory_zero_pad.cpp; 74s local ... (8 actions, 7 running)
ERROR: /private/var/tmp/_bazel_a/f5e9a3325f07a1f02c52d821857db47c/external/org_tensorflow/tensorflow/compiler/xla/BUILD:69:17: ProtoCompile external/org_tensorflow/tensorflow/compiler/xla/xla.pb.h failed: (Illegal instruction): protoc failed: error executing command 
  (cd /private/var/tmp/_bazel_a/f5e9a3325f07a1f02c52d821857db47c/execroot/__main__ && \
  exec env - \
    PATH=/myuserpath/opt/anaconda3/envs/mytraxenv/bin:/myuserpath/opt/anaconda3/condabin:/opt/homebrew/bin:/opt/homebrew/sbin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin:/Library/TeX/texbin:/Library/Apple/usr/bin \
  bazel-out/darwin-opt/bin/external/com_google_protobuf/protoc '--cpp_out=bazel-out/darwin-opt/bin/external/org_tensorflow' -Iexternal/org_tensorflow -Ibazel-out/darwin-opt/bin/external/org_tensorflow -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/any_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/source_context_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/type_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/api_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/descriptor_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/compiler_plugin_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/duration_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/empty_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/field_mask_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/struct_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/timestamp_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/wrappers_proto external/org_tensorflow/tensorflow/compiler/xla/xla.proto)
Execution platform: @local_execution_config_platform//:platform
Target //build:build_wheel failed to build
INFO: Elapsed time: 631.189s, Critical Path: 283.35s
INFO: 2563 processes: 935 internal, 1628 local.
FAILED: Build did NOT complete successfully
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully
b''
Traceback (most recent call last):
  File "/mypath/jax/build/build.py", line 524, in <module>
    main()
  File "/mypath/jax/build/build.py", line 519, in main
    shell(command)
  File "/mypath/jax/build/build.py", line 53, in shell
    output = subprocess.check_output(cmd)
  File "/myuserpath/opt/anaconda3/envs/mytraxenv/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/myuserpath/opt/anaconda3/envs/mytraxenv/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-4.2.1-darwin-x86_64', 'run', '--verbose_failures=true', '--config=avx_posix', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/mypath/jax/dist', '--cpu=x86_64']' returned non-zero exit status 1.

我放弃了构建jaxlib的尝试。在2021年10月底中，发布了一个与m1兼容的jaxlib轮，因此我尝试：

pip install -U pip
pip install -U https://storage.googleapis.com/jax-releases/mac/jaxlib-0.1.74-cp39-none-macosx_11_0_arm64.whl

但得到

ERROR: jaxlib-0.1.74-cp39-none-macosx_11_0_arm64.whl is not a supported wheel on this platform.

尝试了将python从3.9升级到3.10，但得到了同样的信息。

在python中调试，运行

import platform
print(platform.machine())

这表明我的python仍然运行在x86体系结构上，因为Anaconda还没有为M1/ARM提供软件包。

jaxlib [见成功(到目前为止)安装在下面的答案中]

回到trax

在成功安装jaxlib和jax之后，当尝试用(微型锻造的)conda install trax安装trax时，我得到：

Collecting package metadata (current_repodata.json): done
Solving environment: failed with initial frozen solve. Retrying with flexible solve.
Collecting package metadata (repodata.json): done
Solving environment: failed with initial frozen solve. Retrying with flexible solve.

PackagesNotFoundError: The following packages are not available from current channels:

  - trax

Current channels:

  - https://conda.anaconda.org/conda-forge/osx-arm64
  - https://conda.anaconda.org/conda-forge/noarch

To search for alternate channels that may provide the conda package you're
looking for, navigate to

    https://anaconda.org

and use the search bar at the top of the page.

我试着安装为本身意味着，但在构建hp5y时遇到了错误。

ERROR: Failed building wheel for h5py
ERROR: Could not build wheels for h5py, which is required to install pyproject.toml-based projects

现在检查这个错误，并考虑再次使用miniforge，但使用Anaconda的通道安装trax。我到了就会更新。

任何指针或任何设法以不同方式在trax Mac上安装M1的人？

Answer 1

jaxlib + jax

在有人声称成功地使用了微型锻造。之后，我阅读了这，并观看这澄清使用Anaconda和迷你锻造一起。

我用方法安装了微型锻造。由于某种原因，当我运行小型锻造的conda init时，它在~/.bash_profile中设置了初始化代码，尽管我使用的是zsh。我试着手动将代码放在~/.zprofile中，但是它不会加载交互式shell，所以我只是把它放在Anaconda将其初始化代码放在~/zshrc中的地方。

这使得迷你锻造成为默认的管理器。在上面非常有用的视频之后，我创建了一个~/.start_anaconda.sh脚本，这样我就可以使用Anaconda作为替代。

与迷你锻造I

• 使用mytraxenv创建了一个新的conda环境conda create -n mytraxenv python=3，该环境目前有python3.10.2
• 激活环境：conda activate mytraxenv
• 运行conda install numpy和conda install six以确保numpy。six和wheel (由前两者之一安装)安装在我的mytraxenv环境中
• 再试一次，附带一个稍微更新的版本(来自这里)： pip安装-U pip pip安装-U arm64.whl

这在安装jaxlib时起了作用！

然后，我跟随这些指示安装了jax

pip install --upgrade pip
pip install --upgrade "jax[cpu]"

也起作用了。请注意，在python中运行import jax时，它当前警告：

/mytraxenv/lib/python3.10/site-packages/jax/_src/lib/__init__.py:32: 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. "

trax

还没有成功安装trax。

Answer 2

Anaconda源构建

截至2022年2月7日，使用Anaconda安装python失败，因为Anaconda不为python提供臂轮。

当他们提供的时候，从源头构建应该自动工作。如果没有，则克隆存储库，并使用以下方法：

python build/build.py --target_cpu darwin_arm64  

如果bazel不构建并且您正在使用自制，请考虑安装它，并删除下载的bazel文件。

brew install bazel
rm -rf bazel-*

Arm64 Python车轮

我建议从官方的arm64网站下载CPython轮子。可执行文件将在/usr/local/bin/python3中执行，pip3也是如此。

这将允许您使用

/usr/local/bin/pip3 install --upgrade 'jax[cpu]'

若要使用命令行/路径访问正确的python，请将文件夹添加到路径中。

export PATH=/usr/share/bin:$PATH

虚拟环境

我建议不要使用上述方法，而是使用virtual-env (从这里开始使用https://virtualenv.pypa.io )创建虚拟环境(如anaconda )，并在上面的虚拟环境中安装jax。

诗词

诗歌是一个更好的选择，因为它使用虚拟env，并允许您像anaconda一样构建项目和列出依赖关系。

在我编写这篇文章时，有一个小警告:在使用诗歌定义依赖项时安装scipy会导致一个奇怪的错误，使用以下方法手动安装numpy和scipy

poetry shell # to go to virtual environment for the project
pip install --upgrade 'jax[cpu]' jaxlib scipy numpy

特拉克斯

由于anaconda不是选项，所以安装hdf5有点麻烦，但是可以用brew/homebrew来完成。假设您已经激活了正确的venv：

brew install hdf5
export HDF5_DIR=/opt/homebrew/Cellar/hdf5/1.12.1 
pip install trax

不幸的是，这会导致以下错误：

import trax
/Users/quinn/Library/Caches/pypoetry/virtualenvs/no-xNvLksOy-py3.9/lib/python3.9/site-packages/jax/_src/lib/__init__.py:32: 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. "


***************************************************************
Failed to import TensorFlow. Please note that TensorFlow is not installed by default when you install TFDS. This allow you to choose to install either `tf-nightly` or `tensorflow`. Please install the most recent version of TensorFlow, by following instructions at https://tensorflow.org/install.
***************************************************************


Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/quinn/Library/Caches/pypoetry/virtualenvs/no-xNvLksOy-py3.9/lib/python3.9/site-packages/trax/__init__.py", line 18, in <module>
    from trax import layers
  File "/Users/quinn/Library/Caches/pypoetry/virtualenvs/no-xNvLksOy-py3.9/lib/python3.9/site-packages/trax/layers/__init__.py", line 23, in <module>
    from trax.layers.activation_fns import *
  File "/Users/quinn/Library/Caches/pypoetry/virtualenvs/no-xNvLksOy-py3.9/lib/python3.9/site-packages/trax/layers/activation_fns.py", line 26, in <module>
    from trax import math
  File "/Users/quinn/Library/Caches/pypoetry/virtualenvs/no-xNvLksOy-py3.9/lib/python3.9/site-packages/trax/math/__init__.py", line 18, in <module>
    from trax.math import jax as jax_math
  File "/Users/quinn/Library/Caches/pypoetry/virtualenvs/no-xNvLksOy-py3.9/lib/python3.9/site-packages/trax/math/jax.py", line 28, in <module>
    import tensorflow_datasets as tfds
  File "/Users/quinn/Library/Caches/pypoetry/virtualenvs/no-xNvLksOy-py3.9/lib/python3.9/site-packages/tensorflow_datasets/__init__.py", line 43, in <module>
    from tensorflow_datasets.core import tf_compat
  File "/Users/quinn/Library/Caches/pypoetry/virtualenvs/no-xNvLksOy-py3.9/lib/python3.9/site-packages/tensorflow_datasets/core/__init__.py", line 22, in <module>
    tf_compat.ensure_tf_install()
  File "/Users/quinn/Library/Caches/pypoetry/virtualenvs/no-xNvLksOy-py3.9/lib/python3.9/site-packages/tensorflow_datasets/core/tf_compat.py", line 48, in ensure_tf_install
    import tensorflow as tf  # pylint: disable=import-outside-toplevel
ModuleNotFoundError: No module named 'tensorflow'