Bug: Critical Dependency Conflict - JAX 0.4.35 Incompatible with Python 3.12

[alinjar1996]

The current dependency configuration in pyproject.toml specifies requires-python = ">=3.12.0" while simultaneously depending on jax[cuda12]==0.4.35. This version combination is known to be unstable, causing a critical installation failure when Python 3.12 is used.

[vincekurtz]

@alinjar1996 do the changes in #19 fix this problem for you? I'm not able to reproduce the conflict (on Ubuntu 20.04, CUDA version 12.4, NVIDIA driver version 550).

[alinjar1996]

I checked the closed issue #15, and I upgraded JAX to 0.4.38. Although the error message was different. It's working now. It's recommended to update pyproject.toml in main accordingly.