Divergences in Bayesian Non-parametric Causal Inference PyMC 5.10+

Bayesian Non-parametric Causal Inference:
https://www.pymc.io/projects/examples/en/latest/causal_inference/bayesian_nonparametric_causal.html:

Issue description

The notebook was originally developed using pymc 5.3.0, when updating to 5.10 the initial propensity model fit on the logistic regression breaks down. 4000 or so divergences. This seems to be due colinearity and the squared terms in the data set.

*Note that this issue tracker is about the contents in the notebooks,

Expected output

Proposed solution

This issue can be fixed by specifying the init conditions on the sampler.

idata.extend(pm.sample(samples, init='adapt_diag', random_seed=105, idata_kwargs={"log_likelihood": True}))

Start with the linked causal_inference/bayesian_nonparametric_causal.html notebook and its initial propensity model fit, comparing the PyMC 5.3.0 and 5.10 behavior described here. Check the proposed init='adapt_diag' sampler setting and verify that the resulting run no longer produces the reported thousands of divergences and matches the expected output.