[bayes_nonconj] Update lecture

[kp992]

Updates the lecture with:

• Removed pyro and torch code
• Uses only numpyro backed by JAX.
• Typos and Code styling.

Fixes #524

[jstac]

Thanks @kp992 !

[github-actions]

🚀 Deployed on https://68a6f2c4c277859c7aeda77a--nostalgic-wright-5fa355.netlify.app

[jstac]

@mmcky : one for the warning checker: /home/runner/miniconda3/envs/quantecon/lib/python3.12/site-packages/IPython/core/pylabtools.py:170: UserWarning: Creating legend with loc="best" can be slow with large amounts of data. fig.canvas.print_figure(bytes_io, **kw)

[jstac]

@kp992 Great job, much appreciated. Are you willing to hand over to some junior RAs to finish the PR?

Perhaps they can compare the code to the style guide and PEP8, plus add some discussion of the figures at the end. (CC @xuanguang-li @matheusvillasb )

[kp992]

Are you willing to hand over to some junior RAs to finish the PR?

I would quickly send a commit to fix doc reference and then we could hand over for adding more notes.

Perhaps they can compare the code to the style guide and PEP8

I have ran the pep8 and style fixes. One easy way to fix that I use is https://pypi.org/project/jupyter-black/.

import jupyter_black
jupyter_black.load()

And Run All Cells to fix the styling.

[kp992]

Thanks for the review @jstac and @mmcky.

Please feel to merge the PR and happy to hand over this lecture for adding more discussions and notes.

[jstac]

@kp992 thanks again.

@matheusvillasb @xuanguang-li are you willing to take over? Perhaps @mmcky can coordinate merge and opening of a new issue.

[xuanguang-li]

Of course @jstac. I will assign myself to the new issue related to this.

[jstac]

Thanks @xuanguang-li !

[mmcky]

thanks @HumphreyYang -- great guidance.

@xuanguang-li like @HumphreyYang has suggested longer titles are OK when necessary. The style guide is provide guidance around the length (i.e. use the minimum). In this case I agree a longer title is OK to capture the necessary information.

Another idea would be to use

SVI density (uniform prior, normal guide)

sometimes I use this when specifying parameters in a figure for example.

[mmcky]

@xuanguang-li the execution error is due to

/home/runner/_work/lecture-python.myst/lecture-python.myst/lectures/bayes_nonconj.md:877: WARNING: Duplicate explicit target name: "fig_mcmc_uniform".

It looks like there are duplicate labels (they must be unique to work properly).

[xuanguang-li]

Many thanks, @mmcky @HumphreyYang!

[Copilot]

Pull Request Overview

This PR updates the Bayesian non-conjugate lecture to remove dependencies on PyTorch and Pyro, streamlining it to use only NumPyro backed by JAX. The changes include significant code refactoring, typo corrections, and improved code styling throughout the lecture.

Key Changes

• Simplified dependencies: Removed PyTorch and Pyro code, using only NumPyro with JAX backend
• Code refactoring: Updated the BayesianInference class to remove dual solver support and streamline the implementation
• Content improvements: Fixed typos, improved formatting, and enhanced code documentation

[Copilot]

Using jnp.log(0) will result in negative infinity (-inf), which may cause numerical issues. Consider using a small positive value instead, such as jnp.log(1e-10) or define an appropriate lower bound.

Suggested change

	low=jnp.log(0), high=jnp.log(1), loc=loc, scale=scale
	low=jnp.log(1e-10), high=jnp.log(1), loc=loc, scale=scale

[Copilot]

The integrand should be p(Y|θ)p(θ) not p(Y|θ)p(Y). The correct formula should be p(Y) = ∫ p(Y|θ)p(θ) dθ.

Suggested change

	p\left(Y\right)=\int p\left(Y\mid\theta\right)p\left(Y\right) d\theta.
	p\left(Y\right)=\int p\left(Y\mid\theta\right)p\left(\theta\right) d\theta.

[Copilot]

The denominator should be q(θ)p(Y) not p(θ)q(Y). The correct expression should be p(θ,Y)/(q(θ)p(Y)).

Suggested change

	& =-\int q(\theta)\log\frac{p(\theta,Y)}{p(\theta)q(Y)} d\theta\\
	& =-\int q(\theta)\log\frac{p(\theta,Y)}{q(\theta)p(Y)} d\theta\\

[Copilot]

The parameter unpacking is inconsistent. For 'vonMises', kappa is assigned directly from self.param without tuple unpacking, but other distributions use tuple unpacking. This suggests self.param should be a scalar for vonMises, but the documentation shows it as a tuple.

Suggested change

	kappa = self.param
	kappa, = self.param

[Copilot]

The indentation is inconsistent with the surrounding code block. This line should be aligned with the optimizer = nAdam(step_size=lr) line above.

[mmcky]

@HumphreyYang would you mind to review the copilot outputs when you have time.

[xuanguang-li]

Hi @mmcky, thanks for your follow-up. Copilot's suggestions are also very helpful.

I'm currently working on converting the class-based pattern to a JAX pattern and expect to complete it today. Would you prefer I continue with this refactoring, or pause to address some remaining typos first?

[mmcky]

@xuanguang-li please feel free to continue work on this PR. We will merge updates to lectures from 26-August (as there is currently active teaching from this lecture series) so we have a little bit of time.

[xuanguang-li]

Thanks @mmcky. I will push the refactored version later once I resolve those unexpected bugs.

[xuanguang-li]

Hi @mmcky @HumphreyYang,

I attempted to use jax scan to refactor the loop in MCMC_plot, but it didn't work because the length of data[:n] changes dynamically.

After that, I tried lax.dynamic_slice to address this, but the data type is so complex inside the loop that I couldn't convert n into an integer. Now I'm stuck here.

Here is my code:

# plot posteriors
def step(_, n):
    sliced_data = lax.dynamic_slice_in_dim(plot_model.data, 0, n)
    sample_n = MCMC_sampling(
        plot_model.BayesianInferenceClass, sliced_data, num_samples, num_warmup
    )
    return _, sample_n

_, sample = lax.scan(step, None, N_list)

Do you have any idea for this, or would it be better to leave for loops in this lecture?

[jstac]

Thanks for thinking about this so carefully @xuanguang-li !

My view is that we shouldn't be too aggressive in trying to get every speed gain in every situation. For this kind of plotting code, it's fine to do ordinary for loops.

In fact it might be better, because it makes the code easy to read.

For code that's central to algorithms or that will be used many many times, where efficiency is really important, we should push harder and be more aggressive.

[xuanguang-li]

Thanks for your suggestion, @jstac!

It's true that jax scan is difficult to use in plotting functions, as it doesn't seem to support ax.plot inside the loop. I worked around this by separating the data generation and plotting code, but this indeed makes the code less straightforward.

[jstac]

Thanks @xuanguang-li !

@mmcky @HumphreyYang perhaps we should modify the manual to reflect the fact that we're not optimizing everything -- it's a matter of judgement but the basic rule is we should keep things simple and readable for code that uses the model -- e.g., to generate plots -- but isn't internal to the model or the core algorithms that solve it.