Failing to sample from Dirichlet Process Mixture Model using SGLD

[Habush]

Bug Description

I'm trying to reproduce the example from tensorflow-probability titled Fitting Dirichlet Process Mixture Model Using Preconditioned Stochastic Gradient Langevin Dynamics by using blackjax api. I've used the SGLD class. I tried to follow the example notebook that shows how to use SGMCMC in blackjax, but I can't seem to reproduce the results in the tensorflow example.

Steps/Code to Reproduce

I've implemented the blackjax version of the example in a Google Colab notebook located here

Expected Results

Samples from the posterior where there are three clusters

Actual Results

I'm getting nan values for the variables I'm trying to sample.

Versions

BlackJAX 0.8.2

Additional comment

I suspect there are two causes of the issue:
1. My definition of the log probability function
2. the usage of the step size scheduling function

[rlouf]

Thank you for opening an issue. To help me debug, could you do the following:
1 - Look at the output of logprior and loglikelihood with test values for the parameters;
2 - Do the the same for the gradient estimator.
3 - Try with a constant schedule, using a step_size that is very, very small to begin with

[Habush]

Reducing the step_size from 1e-5 to 1e-12 and using a constant schedule, I now get real values instead of nan, so that made a change. However, it fails to converge to the target distribution as the number of inferred clusters is 19 instead of 3

[rlouf]

The step size is probably too small. Moving this to discussions as it does not seem to be a bug. Have you tried the parameters given in the TFP article?

# Learning rates and decay
starter_learning_rate = 1e-6
end_learning_rate = 1e-10
decay_steps = 1e4

# Number of training steps
training_steps = 10000

# Mini-batch size
batch_size = 20

# Sample size for parameter posteriors
sample_size = 100

[junpenglao]

Most likely it is because you are trying to sample bounded variable, which results in non-valid proposal (log prob returns nan).
Recommend to check out https://github.com/blackjax-devs/blackjax/blob/main/examples/change_of_variable_hmc.ipynb

[Habush]

@junpenglao, you are right! I forgot about the constrained support of some of the distributions. After applying the required transformation, I'm able to get the same exact result as the example! (Thanks to @rlouf as well for taking the time to respond)

[rlouf]

That's awesome! Do you plan on making that example public?

[Habush]

If you are open to PR submissions, I can create an example notebook containing the blackjax version, and open a PR. Please let me know.

[rlouf]

By all means please submit a PR,