Fix Long Run Time (bayes_nonconj)

[mmcky]

It is currently taking the cache almost 3 hours to run this lecture set (from scratch).

There are a number of issues reported by tensor flow such as:

/__w/lecture-python.myst/lecture-python.myst/lectures/wealth_dynamics.md: Executing notebook using local CWD [mystnb]
2023-03-22 01:10:05.558920: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65] 
********************************
[Compiling module jit_wealth_time_series_for_loop_jax] Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--=/tmp/foo and attach the results.
********************************

which suggests jax isn't running as it should again.

• once quantecon gpu machine is up and running we should test the lectures and jax on the 3080 to benchmark and diagnose what is going on.

[mmcky]

@HumphreyYang it looks like the bayes_nonconj maybe uninstalling jax due to the following

# install dependencies
!pip install numpyro pyro-ppl torch jax

It looks like torch is reinstalling all the Cuda infrastructure which may mean jax isn't running as it should and could explain the long run time.

Not urgent but would you be able to:

1. Compare local runtime of this lecture to the web version
2. See if we can minimise the installs here so that Cuda libraries are not reinstalled.

[mmcky]

See output

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

[HumphreyYang]

Many thanks @mmcky, I will look into these options.

[HumphreyYang]

Hi @mmcky,

I tested this locally and on Colab. it still runs for around 30 mins. It takes around 1 minute for each MCMC and SVI simulation.

I think it is expected as the size of the simulation is quite large and MCMC is quite slow.

See if we can minimise the installs here so that Cuda libraries are not reinstalled.

Both of them (torch and jax) played important roles in this lecture. I think we can put torch and jax into our requirement file to minimize reinstalls.