I am trying to install Jax, Jaxlib and Chex on Docker, on a Apple M1 Pro machine. The Docker image's base OS is Debian. Doing uname -m gives aarch64 inside the container and arm64 in my local terminal. I am able to install Jax with no issues via pip, but pip install jaxlib raises:
ERROR: Could not find a version that satisfies the requirement jaxlib (from versions: none)
ERROR: No matching distribution found for jaxlib
I tried building it from source following these steps
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
Last command fails:
...
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')
(Log file is empty.) According to this it might be a memory issue so I ran MAX_JOBS=8 python build/build.py which gives something similar but later
...
[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')
(Log file is empty.) I also tried different numbers, which give different errors at different points (either server termination or gcc failure). E.g. MAX_JOBS=5 results in:
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 gives
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
Different runs give different errors...
