Remove PyTorch and Pyro-PPL from workflows#689
Conversation
- Remove torch, torchvision, torchaudio installation from all workflows - Remove pyro-ppl installation from all workflows - Update step names from 'Install JAX, Numpyro, PyTorch' to 'Install JAX and Numpyro' - Disable build cache in ci.yml for full test run - Keep JAX and NumPyro installations
thanks @HumphreyYang |
|
📖 Netlify Preview Ready! Preview URL: https://pr-689--sunny-cactus-210e3e.netlify.app (8c0d609) |
|
📖 Netlify Preview Ready! Preview URL: https://pr-689--sunny-cactus-210e3e.netlify.app (90f6a6f) |
Changes jax[cuda12-local] to jax[cuda12_pip] to avoid cuDNN version compatibility issues. The cuda12_pip variant includes compatible CUDA and cuDNN libraries bundled with JAX, preventing runtime errors from mismatched local CUDA installations. Fixes cuDNN 9.10.0 backward-compatibility error.
|
📖 Netlify Preview Ready! Preview URL: https://pr-689--sunny-cactus-210e3e.netlify.app (f652c69) |
|
@jstac I am going to rebuild the I am currently getting 
so still issues with stability in CUDANN driver versions. Would you like the lectures live for tomorrow? Let me know and I will get up early. |
Updates the quantecon_ubuntu2404 AMI from ami-09baf66e396fa7cfd to ami-0edec81935264b6d3 which includes the latest CUDA and cuDNN libraries for improved compatibility with JAX.
Changes all workflows to use jax[cuda12-local] to leverage the CUDA and cuDNN libraries pre-installed in the new AMI (ami-0edec81935264b6d3). This is faster than cuda12_pip and uses the system libraries. Also removes version pin (==0.6.2) from cache.yml and publish.yml to make all workflows consistent.
Changes all workflows from cuda12-local to cuda13-local to match the CUDA 13 installation in the new AMI (ami-0edec81935264b6d3).
Changes from 'jax[cuda13-local]' to 'jax[cuda13]' following the official JAX documentation for CUDA 13 support. Also uses -U flag instead of --upgrade for consistency with JAX docs.
|
📖 Netlify Preview Ready! Preview URL: https://pr-689--sunny-cactus-210e3e.netlify.app (4d9f83b) |
|
🥳 |
|
@jstac merging this and will rebase your PRs. |
|
📖 Netlify Preview Ready! Preview URL: https://pr-689--sunny-cactus-210e3e.netlify.app (69ba280) |
|
📖 Netlify Preview Ready! Preview URL: https://pr-689--sunny-cactus-210e3e.netlify.app (77f0c62) |
2 participants
This PR removes PyTorch and Pyro-PPL dependencies from all workflows and updates JAX to use CUDA 13 with the new AMI.
Changes Made:
Dependency Removal
JAX Configuration
jax[cuda13]for CUDA 13.0 supportpython scripts/test-jax-install.pyto verify JAX installationInfrastructure Updates
.github/runs-on.ymlwith new AMI:ami-0edec81935264b6d3us-west-2Workflow Improvements
Technical Details:
Dependencies:
This PR requires PR #693 to be merged first, as it updates the runs-on.yml configuration on main (RunsOn reads this from the main branch).