在
返回
trax
trax新手，我正在尝试在本地运行它(macOS 12.1，Apple硅ARM M1处理器，8GB RAM，Anaconda)，但是我遇到了一些问题。
在使用python 3.8.5的环境中，我在(Anaconda)conda环境中安装了运行pip3 install trax==1.3.9trax。后来，我在尝试使用from trax import layers as tl在代码中导入trax层时遇到了问题
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
然后我用python 3.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.
在trax之前安装jaxlib、jax
然后我尝试按照这些说明从源代码构建 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[0mnnBazel binaries for all official releases can be downloaded from here:n https://github.com/bazelbuild/bazel/releasesnnYou 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[0mnnBazel binaries for all official releases can be downloaded from here:n https://github.com/bazelbuild/bazel/releasesnnYou 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后，当尝试使用(miniforge的)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.
我尝试按照trax本身的建议进行安装，但遇到了构建错误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.当我到达它时，我会更新。
任何指针或任何设法通过任何不同方式在 M1 Mac 上安装trax的人？
jaxlib+jax
在有人声称成功使用迷你锻造后，我读了这篇文章并观看了这篇文章，以澄清将蟒蛇和迷你锻造一起使用。
我用苹果的arm64安装了miniforge：苹果硅方法。出于某种原因，当我运行miniforge的conda init即使我使用的是zsh shell，它也在~/.bash_profile中设置了初始化代码。我尝试手动将代码放入~/.zprofile但它不会加载到交互式 shell 上，所以我最终将其放在 Anaconda 将其初始化代码放在~/zshrc中的位置。
这使得 miniforge 成为默认管理器。按照上面非常有用的视频，我创建了一个~/.start_anaconda.sh脚本，以便我可以使用 Anaconda 作为替代方案。
与迷你锻造我
-
创建了一个新的 conda 环境
mytraxenv，目前具有 Python 3.10.2 的conda create -n mytraxenv python=3 -
激活环境：
conda activate mytraxenv -
跑了
conda install numpy，conda install six确保numpy。six和wheel(由前两个之一安装)已安装在我的mytraxenv环境中 -
再次尝试，略有更新的版本(从这里)：
点安装 -U 点 点安装 -U https://storage.googleapis.com/jax-releases/mac/jaxlib-0.1.75-cp310-none-macosx_11_0_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。
Anaconda Source build
截至 2022 年 2 月 7 日星期一，使用 Python 的 Anaconda 安装安装失败，因为 Anaconda 不提供 Python 的臂轮。
当他们确实提供它们时，从源代码构建应该会自动工作。如果没有，请克隆存储库，并使用以下方法：
python build/build.py --target_cpu darwin_arm64
如果 bazel 未构建并且您正在使用自制软件，请考虑安装它，并删除下载的 bazel 文件。
brew install bazel
rm -rf bazel-*
手臂64 蟒蛇轮
我建议从CPython官方网站下载arm64轮子。 可执行文件将在/usr/local/bin/python3中，pip3也是如此。
这应该允许您使用
/usr/local/bin/pip3 install --upgrade 'jax[cpu]'
要使用命令行/路径访问正确的 python，请将文件夹添加到路径中。
export PATH=/usr/share/bin:$PATH
虚拟环境
我建议不要使用上述方法，而是使用此处 https://virtualenv.pypa.io 的virtual-env，像anaconda一样创建虚拟环境，并在上述虚拟环境中安装jax。
诗歌
诗歌是一个更好的选择，因为它使用虚拟环境，并允许您构建项目并像 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'