Fix bug when model has exactly one variable (#437)

Author: jessegrabowski

pymc_extras/inference/laplace.py

@@ -377,7 +377,10 @@ def sample_laplace_posterior(
     posterior_dist = stats.multivariate_normal(
         mean=mu.data, cov=H_inv, allow_singular=True, seed=rng
     )
+
     posterior_draws = posterior_dist.rvs(size=(chains, draws))
+    if mu.data.shape == (1,):
+        posterior_draws = np.expand_dims(posterior_draws, -1)
 
     if transform_samples:
         constrained_rvs, unconstrained_vector = _unconstrained_vector_to_constrained_rvs(model)

tests/test_laplace.py

@@ -263,3 +263,19 @@ def test_fit_laplace(fit_in_unconstrained_space, mode, gradient_backend: Gradien
         else:
             assert idata.fit.rows.values.tolist() == ["mu", "sigma"]
             np.testing.assert_allclose(idata.fit.mean_vector.values, np.array([3.0, 1.5]), atol=0.1)
+
+
+def test_laplace_scalar():
+    # Example model from Statistical Rethinking
+    data = np.array([0, 0, 0, 1, 1, 1, 1, 1, 1])
+
+    with pm.Model():
+        p = pm.Uniform("p", 0, 1)
+        w = pm.Binomial("w", n=len(data), p=p, observed=data.sum())
+
+        idata_laplace = pmx.fit_laplace(progressbar=False)
+
+    assert idata_laplace.fit.mean_vector.shape == (1,)
+    assert idata_laplace.fit.covariance_matrix.shape == (1, 1)
+
+    np.testing.assert_allclose(idata_laplace.fit.mean_vector.values.item(), data.mean(), atol=0.1)