jax cuda installation issue with latest jax=0.6.0 / 0.6.1

[akashsharma02]

Maybe I just got unlucky, but there seems to be a cuda compatibility issue with recent versions of JAX specifically >=0.6.0

See here:
jax-ml/jax#29042

It maybe good to mention a upper limit on jax version so that things work smoothly! :)

[HussainAther]

Thanks for reporting this! I ran into a similar issue with JAX ≥0.6.0 and CUDA recently as well. It might help to pin the version in requirements.txt or add a note in the README under installation instructions. Happy to submit a quick PR if helpful!

[tc2882]

Hi @akashsharma02 , what is the version of jax and brax should I use?

I got the follow error:

🚀Creating environment
/home/linux/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/brax/io/mjcf.py:480: UserWarning: Brax System, piplines and environments are not actively being maintained. Please see MJX for a well maintained JAX-based physics engine: https://github.com/google-deepmind/mujoco/tree/main/mjx. For a host of environments that use MJX, see: https://github.com/google-deepmind/mujoco_playground.
  warnings.warn(
[0.5904899  0.6560999  0.7289999  0.80999994 0.9        1.        ]
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/linux/anaconda3/envs/dial-mpc/bin/dial-mpc-plan", line 33, in <module>
    sys.exit(load_entry_point('dial-mpc', 'console_scripts', 'dial-mpc-plan')())
  File "/home/linux/dial-mpc/dial_mpc/deploy/dial_plan.py", line 288, in main
    mbd_publisher.main_loop()
  File "/home/linux/dial-mpc/dial_mpc/deploy/dial_plan.py", line 165, in main_loop
    state = self.init_mjx_state(
  File "/home/linux/dial-mpc/dial_mpc/deploy/dial_plan.py", line 143, in init_mjx_state
    pipeline_state = self.pipeline_init_jit(self.env.sys, q, qd)
  File "/home/linux/dial-mpc/dial_mpc/deploy/dial_plan.py", line 60, in pipeline_init
    data = _reformat_contact(sys, data)
  File "/home/linux/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/brax/mjx/pipeline.py", line 32, in _reformat_contact
    elasticity = jp.zeros(contact.pos.shape[0])
AttributeError: 'Data' object has no attribute 'pos'. Did you mean: 'qpos'?
/home/linux/anaconda3/envs/dial-mpc/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 6 leaked shared_memory objects to clean up at shutdown

Thank you very much!