Fail to install jax-neuronx on trn1 (Amazon Linux 2023), no distribution for jax==0.4.31

[slitvinov]

While attempting to install jax-neuronx and the associated dependencies on a trn1.2xlarge with Amazon Linux 2023, I encountered deviations from official documentation, and ultimately failed to install jax==0.4.31.

I followed the instructions at:
https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/jax/setup/jax-setup.html

[neuron]
name=Neuron YUM Repository
baseurl=https://yum.repos.neuron.amazonaws.com
enabled=1
metadata_expire=0
EOF
sudo rpm --import https://yum.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB
sudo yum update -y
sudo yum install kernel-devel-$(uname -r) kernel-headers-$(uname -r) -y
sudo yum install git -y
sudo yum install aws-neuronx-dkms-2.* -y
sudo yum install aws-neuronx-collectives-2.* -y
sudo yum install aws-neuronx-runtime-lib-2.* -y
sudo yum install aws-neuronx-tools-2.* -y
export PATH=/opt/aws/neuron/bin:$PATH

I confirmed the accelerator is visible:

neuron-ls
instance-type: trn1.2xlarge
instance-id: i-03a73ac140834eb76
+--------+--------+--------+--------------+
| NEURON | NEURON | NEURON |     PCI      |
| DEVICE | CORES  | MEMORY |     BDF      |
+--------+--------+--------+--------------+
| 0      | 2      | 32 GB  | 0000:00:1e.0 |
+--------+--------+--------+--------------+

The instruction says to run:

python3 -m pip
/usr/bin/python3: No module named pip

But I first had to install pip:

sudo yum install python3-pip -y

When I tried the recommended version combo, it fails

$ python3 -m pip install jax==0.4.31 jaxlib==0.4.31 jax-neuronx libneuronxla neuronx-cc==2.* --extra-index-url=https://pip.repos.neuron.amazonaws.com
<=2.* --extra-index-url=https://pip.repos.neuron.amazonaws.com
Defaulting to user installation because normal site-packages is not writeable
Looking in indexes: https://pypi.org/simple, https://pip.repos.neuron.amazonaws.com
ERROR: Could not find a version that satisfies the requirement jax==0.4.31 (from versions: 0.0, 0.1, 0.1.1, 0.1.2, 0.1.3, 0.1.4, 0.1.5, 0.1.6, 0.1.7, 0.1.8, 0.1.9, 0.1.10, 0.1.11, 0.1.12, 0.1.13, 0.1.14, 0.1.15, 0.1.16, 0.1.18, 0.1.19, 0.1.20, 0.1.21, 0.1.22, 0.1.23, 0.1.24, 0.1.25, 0.1.26, 0.1.27, 0.1.28, 0.1.29, 0.1.30, 0.1.31, 0.1.32, 0.1.33, 0.1.34, 0.1.35, 0.1.36, 0.1.37, 0.1.38, 0.1.39, 0.1.40, 0.1.41, 0.1.42, 0.1.43, 0.1.44, 0.1.45, 0.1.46, 0.1.47, 0.1.48, 0.1.49, 0.1.50, 0.1.51, 0.1.52, 0.1.53, 0.1.54, 0.1.55, 0.1.56, 0.1.57, 0.1.58, 0.1.59, 0.1.60, 0.1.61, 0.1.62, 0.1.63, 0.1.64, 0.1.65, 0.1.66, 0.1.67, 0.1.68, 0.1.69, 0.1.70, 0.1.71, 0.1.72, 0.1.73, 0.1.74, 0.1.75, 0.1.76, 0.1.77, 0.2.0, 0.2.1, 0.2.2, 0.2.3, 0.2.4, 0.2.5, 0.2.6, 0.2.7, 0.2.8, 0.2.9, 0.2.10, 0.2.11, 0.2.12, 0.2.13, 0.2.14, 0.2.15, 0.2.16, 0.2.17, 0.2.18, 0.2.19, 0.2.20, 0.2.21, 0.2.22, 0.2.23, 0.2.24, 0.2.25, 0.2.26, 0.2.27, 0.2.28, 0.3.0, 0.3.1, 0.3.2, 0.3.3, 0.3.4, 0.3.5, 0.3.6, 0.3.7, 0.3.8, 0.3.9, 0.3.10, 0.3.11, 0.3.12, 0.3.13, 0.3.14, 0.3.15, 0.3.16, 0.3.17, 0.3.18, 0.3.19, 0.3.20, 0.3.21, 0.3.22, 0.3.23, 0.3.24, 0.3.25, 0.4.0, 0.4.1, 0.4.2, 0.4.3, 0.4.4, 0.4.5, 0.4.6, 0.4.7, 0.4.8, 0.4.9, 0.4.10, 0.4.11, 0.4.12, 0.4.13, 0.4.14, 0.4.15, 0.4.16, 0.4.17, 0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23, 0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28, 0.4.29, 0.4.30)
ERROR: No matching distribution found for jax==0.4.31

python3 --version
Python 3.9.23

[slitvinov]

It also does not work with Python 3.11:

sudo yum install -qy python3.11 python3.11-pip
python3.11 -m pip install 'jax-neuronx[stable]' --index-url=https://pip.repos.neuron.amazonaws.com
Defaulting to user installation because normal site-packages is not writeable
Looking in indexes: https://pip.repos.neuron.amazonaws.com
Collecting jax-neuronx[stable]
  Using cached https://pip.repos.neuron.amazonaws.com/jax-neuronx/jax_neuronx-0.6.1.1.0.3499%2B2edccbed-py3-none-any.whl (10 kB)
Collecting libneuronxla>=2.2.8201.0
  Using cached https://pip.repos.neuron.amazonaws.com/libneuronxla/libneuronxla-2.2.8201.0%2Bf46ac1ef-py3-none-linux_x86_64.whl (56.2 MB)
Collecting jax-neuronx[stable]
  Using cached https://pip.repos.neuron.amazonaws.com/jax-neuronx/jax_neuronx-0.6.0.1.0.1296%2B1f770067-py3-none-any.whl (10 kB)
  Using cached https://pip.repos.neuron.amazonaws.com/jax-neuronx/jax_neuronx-0.5.3.1.0.719%2B1d9c17be-py3-none-any.whl (6.9 kB)
  Using cached https://pip.repos.neuron.amazonaws.com/jax-neuronx/jax_neuronx-0.1.3-py3-none-any.whl (9.8 kB)
  Using cached https://pip.repos.neuron.amazonaws.com/jax-neuronx/jax_neuronx-0.1.2-py3-none-any.whl (9.5 kB)
  Using cached https://pip.repos.neuron.amazonaws.com/jax-neuronx/jax_neuronx-0.1.1-py3-none-any.whl (6.1 kB)
ERROR: Cannot install jax-neuronx[stable]==0.1.1, jax-neuronx[stable]==0.1.2, jax-neuronx[stable]==0.1.3, jax-neuronx[stable]==0.5.3.1.0.719+1d9c17be, jax-neuronx[stable]==0.6.0.1.0.1296+1f770067 and jax-neuronx[stable]==0.6.1.1.0.3499+2edccbed because these package versions have conflicting dependencies.

The conflict is caused by:
    jax-neuronx[stable] 0.6.1.1.0.3499+2edccbed depends on jax<=0.6.1 and >=0.4.30
    jax-neuronx[stable] 0.6.0.1.0.1296+1f770067 depends on jax<=0.6.0 and >=0.4.30
    jax-neuronx[stable] 0.5.3.1.0.719+1d9c17be depends on jax<=0.5.3 and >=0.4.14
    jax-neuronx[stable] 0.1.3 depends on jax<=0.5.0 and >=0.4.14
    jax-neuronx[stable] 0.1.2 depends on jax<=0.4.35 and >=0.4.14
    jax-neuronx[stable] 0.1.1 depends on jax<=0.4.31 and >=0.4.14

To fix this you could try to:
1. loosen the range of package versions you've specified
2. remove package versions to allow pip attempt to solve the dependency conflict

ERROR: ResolutionImpossible: for help visit https://pip.pypa.io/en/latest/topics/dependency-resolution/#dealing-with-dependency-conflicts
WARNING: There was an error checking the latest version of pip.

[mjkoop]

You do need to have Python 3.10+ installed, however, as you noted there is still a problem with 3.11 following those specific steps on AL2023. The second issue is that the --index-url means that pip looks only at that repository for packages, it will fail if there are any required dependencies that are not yet installed. So the error you are seeing is that jax is not installed so none of the jax_neuronx versions can work.

It is possible to install packages using pip with the --extra-index-url instead of --index-url, which means that packages may be installed from the main PyPi repo OR the one you are specifying. This would simplify your installation, however, it is generally not recommended if you want to make sure you are getting packages and associated dependencies from a specific repository (as we are here).

If you want to get the latest stable setup, the following should first ensure you have the right dependencies for jax-neuronx and neuronx-cc.

python3 -m pip install jax==0.6.0 jaxlib==0.6.0 networkx==2.6 scipy==1.12.0 boto3 python-daemon requests-unixsocket psutil islpy==2023.1 protobuf

Then if you install from the Neuron repo it should work:

python3 -m pip install jax-neuronx[stable] libneuronxla neuronx-cc --index-url=https://pip.repos.neuron.amazonaws.com

It would be good for the team to update the docs to better show this scenario.

[slitvinov]

@mjkoop Thank you.

python3.11 -m pip install jax==0.6.0 jaxlib==0.6.0 networkx==2.6 scipy==1.12.0 boto3 python-daemon requests-unixsocket psutil islpy==2023.1 protobuf

followed by

python3.11 -m pip install jax-neuronx[stable] libneuronxla neuronx-cc --extra-index-url=https://pip.repos.neuron.amazonaws.com

installs the packages.

I have a warning

WARNING:2025-08-15 15:35:36,969:jax._src.xla_bridge:940: Platform 'neuron' is experimental and not all JAX functionality may be correctly supported!
2025-08-15 15:35:42.344297: W neuron/nrt_adaptor.cc:53] nrt_tensor_write_hugepage() is not available, will fall back to nrt_tensor_write().
2025-08-15 15:35:42.344320: W neuron/nrt_adaptor.cc:62] nrt_tensor_read_hugepage() is not available, will fall back to nrt_tensor_read().
2025-Aug-15 15:35:42.0344 3626:3656 [1] net_plugin.cc:73 CCOM WARN NET/Plugin: Error: libnccom-net.so load failed. libfabric.so.1: cannot open shared object file: No such file or directory. Please make sure to install the latest version of libfabric from https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz.
.Completed run_backend_driver.

Compiler status PASS
1

These change to different warnings if I follow the instructions:

curl -O https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz
tar -xzf aws-efa-installer-latest.tar.gz
cd aws-efa-installer-latest
sudo ./efa_installer.sh -y
ldconfig -p | grep libfabric
	libfabric.so.1 (libc6,x86-64) => /opt/amazon/efa/lib64/libfabric.so.1
	libfabric.so (libc6,x86-64) => /opt/amazon/efa/lib64/libfabric.so

python3.11 -c 'import jax; print(jax.numpy.multiply(1, 1))'
WARNING:2025-08-15 15:43:38,438:jax._src.xla_bridge:940: Platform 'neuron' is experimental and not all JAX functionality may be correctly supported!
2025-08-15 15:43:44.044525: W neuron/nrt_adaptor.cc:53] nrt_tensor_write_hugepage() is not available, will fall back to nrt_tensor_write().
2025-08-15 15:43:44.044549: W neuron/nrt_adaptor.cc:62] nrt_tensor_read_hugepage() is not available, will fall back to nrt_tensor_read().
2025-Aug-15 15:43:44.0046 72025:72050 [0] int nccl_net_ofi_create_plugin(nccl_net_ofi_plugin_t**):213 CCOM WARN NET/OFI Failed to initialize sendrecv protocol
2025-Aug-15 15:43:44.0055 72025:72050 [0] int nccl_net_ofi_create_plugin(nccl_net_ofi_plugin_t**):354 CCOM WARN NET/OFI aws-ofi-nccl initialization failed
2025-Aug-15 15:43:44.0065 72025:72050 [0] ncclResult_t nccl_net_ofi_init_no_atexit_fini_v6(ncclDebugLogger_t):183 CCOM WARN NET/OFI Initializing plugin failed
2025-Aug-15 15:43:44.0075 72025:72050 [0] net_plugin.cc:97 CCOM WARN OFI plugin initNet() failed is EFA enabled?
.Completed run_backend_driver.

Compiler status PASS
1

[apoorvtintin]

2025-08-15 15:35:42.344297: W neuron/nrt_adaptor.cc:53] nrt_tensor_write_hugepage() is not available, will fall back to nrt_tensor_write().
2025-08-15 15:35:42.344320: W neuron/nrt_adaptor.cc:62] nrt_tensor_read_hugepage() is not available, will fall back to nrt_tensor_read().

These warnings are harmless and should not impact performance.

2025-Aug-15 15:43:44.0046 72025:72050 [0] int nccl_net_ofi_create_plugin(nccl_net_ofi_plugin_t**):213 CCOM WARN NET/OFI Failed to initialize sendrecv protocol
2025-Aug-15 15:43:44.0055 72025:72050 [0] int nccl_net_ofi_create_plugin(nccl_net_ofi_plugin_t**):354 CCOM WARN NET/OFI aws-ofi-nccl initialization failed
2025-Aug-15 15:43:44.0065 72025:72050 [0] ncclResult_t nccl_net_ofi_init_no_atexit_fini_v6(ncclDebugLogger_t):183 CCOM WARN NET/OFI Initializing plugin failed
2025-Aug-15 15:43:44.0075 72025:72050 [0] net_plugin.cc:97 CCOM WARN OFI plugin initNet() failed is EFA enabled?

The EFA warnings are harmless as EFA is not available in trn1.2xlarge instance-type as mentioned on the Trainium 1 product page. The warning should go away on trn1.32xlarge and trn1n.32xlarge.

[truongp-aws]

Hello,
Thanks for bringing this issue to our attention. JAX requires Python 3.10 or newer. To upgrade Python on AL2023, please refer to this AWS guide. Our documentation has been updated to show setup steps for newer Python versions on AL2023 and will be pusblished in our upcoming NeuronSDK release.

Thanks!