Corrected Python package configuration files.

[romanwixinger]

Problem:

• Before, there was some redundancy and the package numbers of jax and jaxlib were not specified.
• This led to issues in the installation process.

Solution:

• We removed the redundancy and fixed the versions.