Changed the requirements to also work on Mac OS (CPU only)

[kristof]

Platform-Specific Dependency Marker

    "jax[cuda12]>=0.5.0; sys_platform == 'linux'",

This line uses a PEP 508 environment marker (sys_platform == 'linux'). This conditional dependency means:

• On Linux: Both jax>=0.5.0 AND jax[cuda12]>=0.5.0 are installed (the latter brings in CUDA 12 GPU support)
• On Mac OS: Only jax>=0.5.0 (line 9) is installed — the CUDA extra is skipped entirely

This is critical because:

1. CUDA is NVIDIA-only — Mac OS doesn't support CUDA (Apple uses its own Metal GPU framework)
2. Without this marker, pip would try to install jax[cuda12] on Mac, which would fail due to missing CUDA libraries

Summary

The setup is smart: it provides GPU acceleration on Linux via CUDA while gracefully falling back to CPU-only (or Metal) execution on Mac. The sys_platform == 'linux' marker is what makes this "just work" across platforms without requiring separate requirement files.