Support for Jax > 0.4.36?

[RileySE]

Hi, is there a plan to add support for more recent jax versions, especially ones that support cuda 13 or greater? We are trying to use JaxMARL for a project, but currently JaxMARL seems to be unusable with more recent GPU drivers due to breaking API changes in more recent jax versions that support cuda 13.

[amacrutherford]

This is planned but unsure on timeframe. Using the Dockerfile is the best way to go for compatibility currently