jax: Building error on ARM Device

So I’m trying to build JAX from source for the NVIDIA Jetson TX2 which is an ARM device.

I’ve got bazel up and running and it almost builds up until the following error:

[9,973 / 10,056] Compiling external/llvm-project/llvm/lib/CodeGen/InlineSpiller.cpp; 19s local ... (4 actions running)
Target //build:install_xla_in_source_tree up-to-date:
  bazel-bin/build/install_xla_in_source_tree
INFO: Elapsed time: 16931.191s, Critical Path: 374.05s
INFO: 6229 processes: 6229 local.
INFO: Build completed successfully, 10379 total actions
INFO: Running command line: bazel-bin/build/install_xla_in_source_tree /home/nvidia/jax/build
INFO: Build completed successfully, 10379 total actions
Traceback (most recent call last):
  File "/home/nvidia/.cache/bazel/_bazel_nvidia/a5643b5cc286b9b13a96818003a4a7dd/execroot/__main__/bazel-out/arm-opt/bin/build/install_xla_in_source_tree.runfiles/__main__/build/install_xla_in_source_tree.py", line 94, in <module>
    copy(r.Rlocation("__main__/jaxlib/cusolver_kernels.so"))
  File "/home/nvidia/.cache/bazel/_bazel_nvidia/a5643b5cc286b9b13a96818003a4a7dd/execroot/__main__/bazel-out/arm-opt/bin/build/install_xla_in_source_tree.runfiles/__main__/build/install_xla_in_source_tree.py", line 54, in copy
    _copy_so(src_file, dst_dir, dst_filename=dst_filename)
  File "/home/nvidia/.cache/bazel/_bazel_nvidia/a5643b5cc286b9b13a96818003a4a7dd/execroot/__main__/bazel-out/arm-opt/bin/build/install_xla_in_source_tree.runfiles/__main__/build/install_xla_in_source_tree.py", line 43, in _copy_so
    shutil.copy(src_file, dst_file)
  File "/usr/lib/python3.6/shutil.py", line 245, in copy
    copyfile(src, dst, follow_symlinks=follow_symlinks)
  File "/usr/lib/python3.6/shutil.py", line 121, in copyfile
    with open(dst, 'wb') as fdst:
PermissionError: [Errno 13] Permission denied: '/home/nvidia/jax/build/jaxlib/cusolver_kernels.so'
Traceback (most recent call last):
  File "build/build.py", line 457, in <module>
    main()
  File "build/build.py", line 452, in main
    shell(command)
  File "build/build.py", line 51, in shell
    output = subprocess.check_output(cmd)
  File "/usr/lib/python3.6/subprocess.py", line 356, in check_output
    **kwargs).stdout
  File "/usr/lib/python3.6/subprocess.py", line 438, in run
    output=stdout, stderr=stderr)
subprocess.CalledProcessError: Command '['bazelisk', 'run', '--verbose_failures=true', '--config=short_logs', '--config=mkl_open_source_only', '--config=cuda', '--define=xla_python_enable_gpu=true', ':install_xla_in_source_tree', '/home/nvidia/jax/build', '--cpu=arm']' returned non-zero exit status 1.

Looking more closely, I notice I get a permission error thrown at me for this path: PermissionError: [Errno 13] Permission denied: '/home/nvidia/jax/build/jaxlib/cusolver_kernels.so'

I’ve tried modifying the permission with chmod and chown but with no luck.

I’d really appreciate if anyone had any pointers or guidance to resolving this, thank you!