[BUG] LLVM ERROR: mma16816 data type not supported

[UsaidPro]

Describe the bug

Hello! I am trying to get Mava working to test out the library. Following the README, I created a 3.9.10 virtualenv and installed jax[cuda12_local] (I have existing CUDA 12.3 install) and pip install -e . in the mava/ Github root directory.

When I try to run python mava/systems/ppo/ff_ippo.py, I get the following failure:

/home/usaidm/Documents/ExternalLibs/mava/mava/systems/ppo/ff_ippo.py:410: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
  env_states = jax.tree_map(reshape_states, env_states)
/home/usaidm/Documents/ExternalLibs/mava/mava/systems/ppo/ff_ippo.py:411: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
  timesteps = jax.tree_map(reshape_states, timesteps)
/home/usaidm/Documents/ExternalLibs/mava/mava/systems/ppo/ff_ippo.py:431: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
  replicate_learner = jax.tree_map(broadcast, replicate_learner)
[... (default config is printed here)]

LLVM ERROR: mma16816 data type not supported
Aborted (core dumped)

To Reproduce

Steps to reproduce the behavior:

1. pip install "jax[cuda12_local]"
2. Follow the Python virtual environment-based installation guide in DETAILED_INSTALL.md
3. Run python mava/systems/ppo/ff_ippo.py in mava/ Github root directory

Expected behavior

Successfully start Mava and train on RobotWarehouse env.

Context (Environment)

• Mava (and other package versions): result of pip freeze > requirements.txt - Github gist link here
• OS: Ubuntu 22.04
• NVIDIA driver: 545.23.08
• CUDA version: 12.1 (I have verified this works with Jax)
• CuDNN version: 9.2.0 (I don't know how to confirm this)
• GPU: RTX 4090

Additional context

N/A

Possible Solution

Not sure since I am new to package. I will test other environments to see if they also have the same issue.