LECTURE: An Introduction to JAX (#237)

* Added jax lecture

* disable build cache, enable execution reports upon failure

* enable upload of build failure reports for all outputs

* catch error using raises-exception tag

* Fix missing variables for solution

* enable html build only

* fix comment typo

* Add jax to environment

* add jaxlib and jax[cuda]

* Update lectures/jax_intro.md

Co-authored-by: Smit Lunagariya <[email protected]>

* Install jax outside of environment creation

* update bash and remove pip upgrade

* move to install jax from conda-forge

* Remove sandpit packages and revert to pip

* Allow 4 minute timeout for cell execution

* enable 10min timeout

* Use jax install from Docs

* Install nvidia drivers to check if driver issue

* use nvidia cuda docker container

* remove nvidia driver install step

* disable latex support for now

* try install cudatoolkit

* give container access to gpu

* remove locate and keep whereis

* Upgrade CUDA to latest

* upgrade cuda before installing jax

* Max HDD bigger for new cuda drivers

* use latest nvidia docker container release using 11.8

* match cuda toolkit from conda

* ensure cuda installed as reported not available

* enable nvidia-smi

* add diagnostics

* add nvidia runtime to docker launch

* enable nvidia-smi on jb run

* revert to cml base image

* try installing within the container

* Test quantecon local runner

* simply and see if can run on local conda env

* fix syntax for conda

* run containerised version

* enable full preview build including pdf, download notebook + build cache

* remove install of latex as installed locally on runner

* add explanations for block_until_ready()

* change wording

* improve based on comments

* switch back to using ec2

* remove cache

* make sure nvcc binaries are available

* check solution from jax issues

* install cuda and cuda-toolkit

* try using python in build to container

* add python packages for lectures

* use nvidia as base docker

* move back to conda

* rely on nvidia docker for drivers

* ensure jax is uninstalled

* try nvidia cuda=11.2

* Check tensorflow docker container

* check tensorflow==2.9.2

* change hardware to V100

* remove tensorflow, use nvidia docker (quicker), check opt_savings

* remove tensorflow docker, update title for test of opt_savings

* tidy up, enable pdf and download nb builds

* remove sudo

* setup timezone data

* remove ENV

* remove infrastructure testing and move to new PR

* remove test file from toc

* check conda environment

* more conda path debug

* remove debu

* add hardware details

* add more useful information in lecture

* improve note in lecture

Co-authored-by: mmcky <[email protected]>
Co-authored-by: mmcky <[email protected]>
Co-authored-by: Smit Lunagariya <[email protected]>
Co-authored-by: Humphrey Yang <[email protected]>

Author: 4 people

lectures/_config.yml

@@ -5,7 +5,7 @@ description: This website presents a set of lectures on python programming for e
 
 execute:
   execute_notebooks: "cache"
-  timeout: 120
+  timeout: 600 # 10 minutes
 #   run_in_temp: true
 
 html:

lectures/_toc.yml

@@ -24,6 +24,7 @@ parts:
   chapters:
   - file: numba
   - file: parallelization
+  - file: jax_intro
 - caption: Advanced Python Programming
   numbered: true
   chapters:
@@ -34,4 +35,4 @@ parts:
   numbered: true
   chapters:
   - file: troubleshooting
-  - file: status
+  - file: status