Installing jax on a Virtual Machine of Ubuntu 22.04 and AVX support isn't working

[AlonzoA2-code]

Be advised that I am a complete newbie when it comes to jax. I am attempting to run trajax in a virtual machine of Ubuntu 22.04 running on a Windows 11 HP laptop in Oracle VM VirtualBox 7.0.8. After installing jaxlib, jax, and trajax according to the instructions from their respective GitHubs, I have run into this error constantly when trying to import jax in python:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[2], line 1
----> 1 from jax.config import config
      2 config.update('jax_enable_x64', True)
      4 import jax.numpy as jnp

File [~/.local/lib/python3.10/site-packages/jax/__init__.py:35](https://file+.vscode-resource.vscode-cdn.net/home/alonzoa2/Documents/trajax-main/notebooks/l4dc/~/.local/lib/python3.10/site-packages/jax/__init__.py:35)
     30 del _cloud_tpu_init
     32 # Confusingly there are two things named "config": the module and the class.
     33 # We want the exported object to be the class, so we first import the module
     34 # to make sure a later import doesn't overwrite the class.
---> 35 from jax import config as _config_module
     36 del _config_module
     38 # Force early import, allowing use of `jax.core` after importing `jax`.

File [~/.local/lib/python3.10/site-packages/jax/config.py:17](https://file+.vscode-resource.vscode-cdn.net/home/alonzoa2/Documents/trajax-main/notebooks/l4dc/~/.local/lib/python3.10/site-packages/jax/config.py:17)
      1 # Copyright 2018 The JAX Authors.
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     14 
     15 # TODO(phawkins): fix users of this alias and delete this file.
---> 17 from jax._src.config import config  # noqa: F401
...
---> 84 cpu_feature_guard.check_cpu_features()
     86 # TODO(phawkins): remove after minimium jaxlib version is 0.4.9 or newer.
     87 try:

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.

I have verified that my host machine has AVX support (my CPU is an Intel Core i5-8250U) and I have turned on the setting that should allow VirtualBox to access my computer's AVX capabilities (Nested VT-x/AMD-V). I also tried building jaxlib from source according to this (https://jax.readthedocs.io/en/latest/developer.html), but at the step to build jaxlib with only CPU support, the process always fails with this error:

alonzoa2@Ubuntu-22:~/jax$ python3 build/build.py
     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\


Bazel binary path: ./bazel-6.1.2-linux-x86_64
Bazel version: 6.1.2
Python binary path: /usr/bin/python3
Python version: 3.10
NumPy version: 1.25.0
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-6.1.2-linux-x86_64 run --verbose_failures=true :build_wheel -- --output_path=/home/alonzoa2/jax/dist --cpu=x86_64

Starting local Bazel server and connecting to it...

INFO: Options provided by the client:

  Inherited 'common' options: --isatty=0 --terminal_columns=80

INFO: Reading rc options for 'run' from /home/alonzoa2/jax/.bazelrc:

  Inherited 'common' options: --experimental_repo_remote_exec

INFO: Reading rc options for 'run' from /home/alonzoa2/jax/.bazelrc:

  Inherited 'build' options: --nocheck_visibility --apple_platform_type=macos --macos_minimum_os=10.14 --announce_rc --define open_source_build=true --spawn_strategy=standalone --enable_platform_specific_config --experimental_cc_shared_library --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 --define=tsl_link_protobuf=true -c opt --config=short_logs --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir. --@xla//xla/python:enable_gpu=false --@xla//xla/python:enable_tpu=false

INFO: Reading rc options for 'run' from /home/alonzoa2/jax/.jax_configure.bazelrc:

  Inherited 'build' options: --strategy=Genrule=standalone --repo_env PYTHON_BIN_PATH=/usr/bin/python3 --action_env=PYENV_ROOT --python_path=/usr/bin/python3 --distinct_host_configuration=false --config=avx_posix --config=mkl_open_source_only

INFO: Found applicable config definition build:short_logs in file /home/alonzoa2/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING

INFO: Found applicable config definition build:avx_posix in file /home/alonzoa2/jax/.bazelrc: --copt=-mavx --host_copt=-mavx

INFO: Found applicable config definition build:mkl_open_source_only in file /home/alonzoa2/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1

INFO: Found applicable config definition build:linux in file /home/alonzoa2/jax/.bazelrc: --config=posix --copt=-Wno-unknown-warning-option --copt=-Wno-stringop-truncation --copt=-Wno-array-parameter

INFO: Found applicable config definition build:posix in file /home/alonzoa2/jax/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++17 --host_cxxopt=-std=c++17

Loading: 

Loading: 

DEBUG: /home/alonzoa2/.cache/bazel/_bazel_alonzoa2/2d724c09cb4334d607a544256bfca8d9/external/xla/third_party/repo.bzl:132:14: 

Warning: skipping import of repository 'tf_runtime' because it already exists.

DEBUG: /home/alonzoa2/.cache/bazel/_bazel_alonzoa2/2d724c09cb4334d607a544256bfca8d9/external/xla/third_party/repo.bzl:132:14: 

Warning: skipping import of repository 'llvm-raw' because it already exists.

Loading: 

Loading: 

Loading: 

Loading: 1 packages loaded

Analyzing: target //build:build_wheel (2 packages loaded, 0 targets configured)

Analyzing: target //build:build_wheel (19 packages loaded, 10 targets configured)

Analyzing: target //build:build_wheel (35 packages loaded, 10 targets configured)

Analyzing: target //build:build_wheel (48 packages loaded, 227 targets configured)

Analyzing: target //build:build_wheel (59 packages loaded, 327 targets configured)

Analyzing: target //build:build_wheel (80 packages loaded, 656 targets configured)

Analyzing: target //build:build_wheel (115 packages loaded, 848 targets configured)

Analyzing: target //build:build_wheel (129 packages loaded, 1299 targets configured)

Analyzing: target //build:build_wheel (145 packages loaded, 1477 targets configured)

Analyzing: target //build:build_wheel (152 packages loaded, 1814 targets configured)

Analyzing: target //build:build_wheel (158 packages loaded, 2668 targets configured)

Analyzing: target //build:build_wheel (165 packages loaded, 5587 targets configured)

Analyzing: target //build:build_wheel (174 packages loaded, 8764 targets configured)

Analyzing: target //build:build_wheel (175 packages loaded, 10232 targets configured)

Analyzing: target //build:build_wheel (183 packages loaded, 11036 targets configured)

Analyzing: target //build:build_wheel (189 packages loaded, 12928 targets configured)

Analyzing: target //build:build_wheel (193 packages loaded, 14647 targets configured)

Analyzing: target //build:build_wheel (197 packages loaded, 15059 targets configured)

INFO: Analyzed target //build:build_wheel (197 packages loaded, 15059 targets configured).

INFO: Found 1 target...

[0 / 1] checking cached actions

[7 / 10] checking cached actions

[206 / 211] [Prepa] ProtoCompile external/xla/xla/service/hlo.pb.h

ERROR: /home/alonzoa2/.cache/bazel/_bazel_alonzoa2/2d724c09cb4334d607a544256bfca8d9/external/xla/xla/BUILD:79:17: ProtoCompile external/xla/xla/xla_data.pb.h failed: (Illegal instruction): protoc failed: error executing command (from target @xla//xla:xla_data_proto_genproto) 

  (cd /home/alonzoa2/.cache/bazel/_bazel_alonzoa2/2d724c09cb4334d607a544256bfca8d9/execroot/__main__ && \

  exec env - \

    PATH=/home/alonzoa2/.local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/snap/bin \

  bazel-out/k8-opt-exec-50AE0418/bin/external/com_google_protobuf/protoc '--cpp_out=bazel-out/k8-opt/bin/external/xla' -Iexternal/xla -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/any_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/source_context_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/type_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/api_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/descriptor_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/compiler_plugin_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/duration_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/empty_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/field_mask_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/struct_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/timestamp_proto -Ibazel-out/k8-opt/bin/external/com_google_protobuf/_virtual_imports/wrappers_proto external/xla/xla/xla_data.proto)

# Configuration: c8ef805acd97db9d9c71f0c883180e1d661f614815683ee4eb170b1054f2ce00

# Execution platform: @local_execution_config_platform//:platform

Target //build:build_wheel failed to build

INFO: Elapsed time: 39.898s, Critical Path: 0.85s

INFO: 3 processes: 3 internal.

FAILED: Build did NOT complete successfully

ERROR: Build failed. Not running target

b''

Traceback (most recent call last):

  File "/home/alonzoa2/jax/build/build.py", line 554, in <module>

    main()

  File "/home/alonzoa2/jax/build/build.py", line 549, in main

    shell(command)

  File "/home/alonzoa2/jax/build/build.py", line 53, in shell

    output = subprocess.check_output(cmd)

  File "/usr/lib/python3.10/subprocess.py", line 420, in check_output

    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,

  File "/usr/lib/python3.10/subprocess.py", line 524, in run

    raise CalledProcessError(retcode, process.args,

subprocess.CalledProcessError: Command '['./bazel-6.1.2-linux-x86_64', 'run', '--verbose_failures=true', ':build_wheel', '--', '--output_path=/home/alonzoa2/jax/dist', '--cpu=x86_64']' returned non-zero exit status 1.

I have been dealing with these errors for several days and no solution I've tried from googling this has helped me. What could be causing these errors and how could I fix them? If necessary, I can provide more information about my system/settings. Any and all help is greatly appreciated.