无法在 Apple M1 和 Docker 上安装 Jaxlib

Question

我正在尝试在 Apple M1 Pro 机器上的 Docker 上安装 Jax、Jaxlib 和 Chex。Docker 镜像的基础操作系统是 Debian。在容器内和我的本地终端中进行操作uname -m。我可以通过 pip 安装 Jax，没有任何问题，但会引发：aarch64arm64pip install jaxlib

ERROR: Could not find a version that satisfies the requirement jaxlib (from versions: none)
ERROR: No matching distribution found for jaxlib

我尝试按照以下步骤 从源代码构建它git clone https://github.com/google/jax cd jax apt-get install g++ python python3-dev pip install numpy six wheel python build/build.py

最后一个命令失败：

...
INFO: Found 1 target...
[0 / 224] [Prepa] Creating source manifest for //build:build_wheel
[119 / 2,863] Compiling src/google/protobuf/compiler/objectivec/objectivec_helpers.cc; 1s local ... (5 actions, 4 running)
[223 / 3,067] Compiling snappy.cc; 0s local ... (5 actions, 4 running)
[437 / 3,333] Compiling llvm/lib/Support/BinaryStreamRef.cpp; 0s local ... (5 actions, 4
running)
[576 / 3,333] Compiling llvm/lib/TableGen/Record.cpp; 1s local ... (5 actions, 4 running)
[856 / 3,678] Compiling llvm/lib/TableGen/Record.cpp; 2s local ... (5 actions, 4 running) [1,227 / 4,047] Compiling llvm/utils/TableGen/GlobalISelEmitter.cpp; 13s local ... (5
actions running)
[1,227 / 4,047] Compiling llvm/utils/TableGen/GlobalISelEmitter.cpp; 44s local ... (5
actions running)

Server terminated abruptly (error code: 14, error message: 'Socket closed', log file:
'/root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/server/jvm.out')

（日志文件是空的。）根据这个，这可能是一个内存问题，所以我跑了MAX_JOBS=8 python build/build.py，它给出了类似的东西，但后来

...
[2,710 / 4,162] Compiling re2/dfa.cc; 0s local ... (5 actions, 4 running)
[2,793 / 4,162] Compiling tensorflow/core/ir/ops.cc; 5s local ... (5 actions, 4 running)
[2,803 / 4,162] Compiling tensorflow/core/ir/ops.cc; 29s local ... (5 actions running)
[2,803 / 4,162] Compiling tensorflow/core/ir/ops.cc; 56s local ... (5 actions running)

Server terminated abruptly (error code: 14, error message: 'Socket closed', log file: '/root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/server/jvm.out')

（日志文件为空。）我还尝试了不同的数字，它们在不同的点给出不同的错误（服务器终止或 gcc 失败）。例如MAX_JOBS=5结果：

ERROR: /root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/external/llvm-project/mlir/BUILD.bazel:7723:11: Compiling mlir/lib/Dialect/Arithmetic/IR/ArithmeticDialect.cpp failed: (Exit 1): gcc failed: error executing command

MAX_JOBS=1给出

ERROR: /root/.cache/bazel/_bazel_root/cfd1b2cc6fe180f3eb424db6004de364/external/llvm-project/mlir/BUILD.bazel:3437:11: Compiling mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp failed: (Exit 1): gcc failed: error executing command

不同的运行会产生不同的错误...