Keras 3.0.4|3.0.5 constrain jax to 0.4

[giotherobot]

Solution to issue cannot be found in the documentation.

• [X] I checked the documentation.

Issue

Trying to create a gpu environment with CUDA 12, keras 3 and jax backend fails.

Tested with:

Environment.yaml

``` name: testenv channels: - conda-forge dependencies: - jax>0.4.23 - jaxlib - cuda-version=12.* - cuda-cupti - cuda-nvcc - keras>=3.0.0 ```

Solver output:

Could not solve for environment specs
The following packages are incompatible
├─ jax >0.4.23  is requested and can be installed;
└─ keras >=3.0.0  is not installable because there are no viable options
   ├─ keras 3.0.2 would require
   │  └─ jax 0.4.20 , which conflicts with any installable versions previously reported;
   └─ keras [3.0.4|3.0.5] would require
      └─ jax 0.4 , which conflicts with any installable versions previously reported.

jax >=0.4.23 is required to obtain jaxlib compiled for cuda 12.

Installed packages

None

Environment info

mamba version : 1.5.1
     active environment : None
            shell level : 0
       user config file : /home/giovannif/.condarc
 populated config files : /home/giovannif/mambaforge/.condarc
                          /home/giovannif/.condarc
          conda version : 4.14.0
    conda-build version : not installed
         python version : 3.10.10.final.0
       virtual packages : __cuda=12.4=0
                          __linux=6.5.0=0
                          __glibc=2.35=0
                          __unix=0=0
                          __archspec=1=x86_64
       base environment : /home/giovannif/mambaforge  (writable)
      conda av data dir : /home/giovannif/mambaforge/etc/conda
  conda av metadata url : None
           channel URLs : https://conda.anaconda.org/conda-forge/linux-64
                          https://conda.anaconda.org/conda-forge/noarch
          package cache : /home/giovannif/mambaforge/pkgs
                          /home/giovannif/.conda/pkgs
       envs directories : /home/giovannif/mambaforge/envs
                          /home/giovannif/.conda/envs
               platform : linux-64
             user-agent : conda/4.14.0 requests/2.28.1 CPython/3.10.10 Linux/6.5.0-26-generic neon/22.04 glibc/2.35
                UID:GID : 1000:1000
             netrc file : /home/giovannif/.netrc
           offline mode : False

[traversaro]

To have a more self-contained example, this fails:

traversaro@IITBMP014LW012:~$ conda create -n kerastest "keras==3.0.5" "jax==0.4.23"
Retrieving notices: ...working... done
Channels:
 - conda-forge
Platform: linux-64
Collecting package metadata (repodata.json): done
Solving environment: failed

LibMambaUnsatisfiableError: Encountered problems while solving:
  - package keras-3.0.5-pyhd8ed1ab_0 has constraint jax 0.4 conflicting with jax-0.4.23-pyhd8ed1ab_0

Could not solve for environment specs
The following packages are incompatible
├─ jax 0.4.23  is requested and can be installed;
└─ keras 3.0.5  is not installable because it requires
   └─ jax 0.4 , which conflicts with any installable versions previously reported.