ENH: include nvidia-smi across lectures

lectures/aiyagari_jax.md

@@ -57,6 +57,12 @@ import jax
 import jax.numpy as jnp
 ```
 
+Let's check the GPU we are running
+
+```{code-cell} ipython3
+!nvidia-smi
+```
+
 We will use 64 bit floats with JAX in order to increase the precision.
 
 ```{code-cell} ipython3

lectures/arellano.md

@@ -79,8 +79,16 @@ import jax
 import jax.numpy as jnp
 ```
 
+Let's check the GPU we are running
+
+```{code-cell} ipython3
+!nvidia-smi
+```
+
+We will use 64 bit floats with JAX in order to increase the precision.
+
 ```{code-cell} ipython3
-jax.config.update('jax_enable_x64', True)
+jax.config.update("jax_enable_x64", True)
 ```
 
 ## Structure

lectures/ifp_egm.md

@@ -48,6 +48,11 @@ from numba import njit, float64
 from numba.experimental import jitclass
 ```
 
+Let's check the GPU we are running
+
+```{code-cell} ipython3
+!nvidia-smi
+```
 
 We use 64 bit floating point numbers for extra precision.
 

lectures/inventory_dynamics.md

@@ -48,14 +48,10 @@ from jax import random, lax
 from collections import namedtuple
 ```
 
-Let’s check the backend used by JAX and the devices available
+Let's check the GPU we are running
 
 ```{code-cell} ipython3
-# Check if JAX is using GPU
-print(f"JAX backend: {jax.devices()[0].platform}")
-
-# Check the devices available for JAX
-print(jax.devices())
+!nvidia-smi
 ```
 
 ## Sample paths

lectures/kesten_processes.md

@@ -62,6 +62,12 @@ import jax.numpy as jnp
 from jax import random
 ```
 
+Let's check the GPU we are running
+
+```{code-cell} ipython3
+!nvidia-smi
+```
+
 ## Kesten processes
 
 ```{index} single: Kesten processes; heavy tails

lectures/newtons_method.md

@@ -35,6 +35,12 @@ import jax.numpy as jnp
 from scipy.optimize import root
 ```
 
+Let's check the GPU we are running
+
+```{code-cell} ipython3
+!nvidia-smi
+```
+
 ## The Equilibrium Problem
 
 In this section we describe the market equilibrium problem we will solve with

lectures/opt_invest.md

@@ -25,7 +25,6 @@ We require the following library to be installed.
 !pip install --upgrade quantecon
 ```
 
-
 A monopolist faces inverse demand
 curve
 
@@ -66,6 +65,12 @@ import jax.numpy as jnp
 import matplotlib.pyplot as plt
 ```
 
+Let's check the GPU we are running
+
+```{code-cell} ipython3
+!nvidia-smi
+```
+
 We will use 64 bit floats with JAX in order to increase the precision.
 
 ```{code-cell} ipython3

lectures/opt_savings.md

@@ -33,6 +33,12 @@ import matplotlib.pyplot as plt
 import time
 ```
 
+Let's check the GPU we are running
+
+```{code-cell} ipython3
+!nvidia-smi
+```
+
 Use 64 bit floats with JAX in order to match NumPy code
 - By default, JAX uses 32-bit datatypes.
 - By default, NumPy uses 64-bit datatypes.

lectures/short_path.md

@@ -30,6 +30,12 @@ import jax.numpy as jnp
 import jax
 ```
 
+Let's check the GPU we are running
+
+```{code-cell} ipython3
+!nvidia-smi
+```
+
 ## Solving for Minimum Cost-to-Go
 
 Let $J(v)$ denote the minimum cost-to-go from node $v$,

lectures/status.md

@@ -28,4 +28,10 @@ You can check the backend used by JAX using:
 import jax
 # Check if JAX is using GPU
 print(f"JAX backend: {jax.devices()[0].platform}")
+```
+
+and the hardware we are running on:
+
+```{code-cell} ipython3
+!nvidia-smi
 ```

lectures/wealth_dynamics.md

@@ -29,7 +29,7 @@ import jax.numpy as jnp
 from collections import namedtuple
 ```
 
-Let's check the hardware we are running on:
+Let's check the GPU we are running
 
 ```{code-cell} ipython3
 !nvidia-smi