Add a cell to check JAX backend

Author: Smit-create

lectures/short_path.md

@@ -12,6 +12,7 @@ kernelspec:
 ---
 
 
+
 # Shortest Paths
 
 
@@ -28,6 +29,17 @@ import jax
 ```
 
 
+
+Let’s check the backend used by JAX and the devices available.
+
+```{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())
+```
+
+
 ## Solving for Minimum Cost-to-Go
 
 Let $J(v)$ denote the minimum cost-to-go from node $v$,
@@ -73,6 +85,7 @@ Q = jnp.array([[inf, 1,   5,   3,   inf, inf, inf],
 ```
 
 
+
 Notice that the cost of staying still (on the principle diagonal) is set to
 
 * `jnp.inf` for non-destination nodes --- moving on is required.
@@ -101,6 +114,7 @@ print("The cost-to-go function is", J)
 ```
 
 
+
 We can further optimize the above code by using [jax.lax.while_loop](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.while_loop.html). The extra acceleration is due to the fact that the entire operation can be optimized by the JAX compiler and launched as a single kernel on the GPU.
 
 ```{code-cell} ipython3
@@ -129,6 +143,7 @@ def cond_fun(values):
 ```
 
 
+
 Let's see the timing for JIT compilation of the functions and runtime results.
 
 ```{code-cell} ipython3
@@ -138,6 +153,7 @@ jax.lax.while_loop(cond_fun, body_fun, init_val=(0, J, False))[1]
 ```
 
 
+
 Now, this runs faster once we have the JIT compiled JAX version of the functions.
 
 ```{code-cell} ipython3
@@ -147,6 +163,7 @@ jax.lax.while_loop(cond_fun, body_fun, init_val=(0, J, False))[1]
 ```
 
 
+
 ```{note}
 Large speed gains while using `jax.lax.while_loop` won't be realized unless the shortest path problem is relatively large.
 ```
@@ -311,6 +328,7 @@ def map_graph_to_distance_matrix(in_file):
 ```
 
 
+
 Let's write a function `compute_cost_to_go` that returns $J$ given any valid $Q$.
 
 ```{code-cell} ipython3
@@ -341,6 +359,7 @@ def compute_cost_to_go(Q):
 ```
 
 
+
 Finally, here's a function that uses the `cost-to-go` function to obtain the
 optimal path (and its cost).
 
@@ -359,13 +378,15 @@ def print_best_path(J, Q):
 ```
 
 
+
 Okay, now we have the necessary functions, let's call them to do the job we were assigned.
 
 ```{code-cell} ipython3
 Q = map_graph_to_distance_matrix('graph.txt')
 ```
 
 
+
 Let's see the timings for jitting the function and runtime results.
 
 ```{code-cell} ipython3
@@ -384,6 +405,7 @@ print_best_path(J, Q)
 ```
 
 
+
 The total cost of the path should agree with $J[0]$ so let's check this.
 
 ```{code-cell} ipython3