Enable RunsOn GPU support for lecture builds

- Add scripts/test-jax-install.py to verify JAX/GPU installation
- Add .github/runs-on.yml with QuantEcon Ubuntu 24.04 AMI configuration
- Update cache.yml to use RunsOn g4dn.2xlarge GPU runner
- Update ci.yml to use RunsOn g4dn.2xlarge GPU runner
- Update publish.yml to use RunsOn g4dn.2xlarge GPU runner
- Install JAX with CUDA 13 support and Numpyro on all workflows
- Add nvidia-smi check to verify GPU availability

This mirrors the setup used in lecture-python.myst repository.

Author: mmcky

.github/runs-on.yml

@@ -0,0 +1,6 @@
+images:
+  quantecon_ubuntu2404:
+    platform: "linux"
+    arch: "x64"
+    ami: "ami-0edec81935264b6d3"
+    region: "us-west-2"

.github/workflows/cache.yml

@@ -6,7 +6,7 @@ on:
   workflow_dispatch:
 jobs:
   cache:
-    runs-on: ubuntu-latest
+    runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=quantecon_ubuntu2404/disk=large"
     steps:
       - uses: actions/checkout@v6
       - name: Setup Anaconda
@@ -18,6 +18,16 @@ jobs:
           python-version: "3.13"
           environment-file: environment.yml
           activate-environment: quantecon
+      - name: Install JAX and Numpyro
+        shell: bash -l {0}
+        run: |
+          pip install -U "jax[cuda13]"
+          pip install numpyro
+          python scripts/test-jax-install.py
+      - name: Check nvidia drivers
+        shell: bash -l {0}
+        run: |
+          nvidia-smi
       - name: Build HTML
         shell: bash -l {0}
         run: |

.github/workflows/ci.yml

@@ -2,7 +2,7 @@ name: Build Project [using jupyter-book]
 on: [pull_request]
 jobs:
   preview:
-    runs-on: ubuntu-latest
+    runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=quantecon_ubuntu2404/disk=large"
     steps:
       - uses: actions/checkout@v6
         with:
@@ -16,6 +16,15 @@ jobs:
           python-version: "3.13"
           environment-file: environment.yml
           activate-environment: quantecon
+      - name: Check nvidia Drivers
+        shell: bash -l {0}
+        run: nvidia-smi
+      - name: Install JAX and Numpyro
+        shell: bash -l {0}
+        run: |
+          pip install -U "jax[cuda13]"
+          pip install numpyro
+          python scripts/test-jax-install.py
       - name: Install latex dependencies
         run: |
           sudo apt-get -qq update

.github/workflows/publish.yml

@@ -6,7 +6,7 @@ on:
 jobs:
   publish:
     if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags')
-    runs-on: ubuntu-latest
+    runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=quantecon_ubuntu2404/disk=large"
     steps:
       - name: Checkout
         uses: actions/checkout@v6
@@ -21,6 +21,16 @@ jobs:
           python-version: "3.13"
           environment-file: environment.yml
           activate-environment: quantecon
+      - name: Install JAX and Numpyro
+        shell: bash -l {0}
+        run: |
+          pip install -U "jax[cuda13]"
+          pip install numpyro
+          python scripts/test-jax-install.py
+      - name: Check nvidia drivers
+        shell: bash -l {0}
+        run: |
+          nvidia-smi
       - name: Install latex dependencies
         run: |
           sudo apt-get -qq update

scripts/test-jax-install.py

@@ -0,0 +1,21 @@
+import jax
+import jax.numpy as jnp
+
+devices = jax.devices()
+print(f"The available devices are: {devices}")
+
+@jax.jit
+def matrix_multiply(a, b):
+    return jnp.dot(a, b)
+
+# Example usage:
+key = jax.random.PRNGKey(0)
+x = jax.random.normal(key, (1000, 1000))
+y = jax.random.normal(key, (1000, 1000))
+z = matrix_multiply(x, y)
+
+# Now the function is JIT compiled and will likely run on GPU (if available)
+print(z)
+
+devices = jax.devices()
+print(f"The available devices are: {devices}")