add test for start and start_sigma plus minor fixes

Author: markusschmaus

pymc/tests/test_variational_inference.py

@@ -571,6 +571,28 @@ def test_fit_oo(inference, fit_kwargs, simple_model_data):
     np.testing.assert_allclose(np.std(trace.posterior["mu"]), np.sqrt(1.0 / d), rtol=0.2)
 
 
+def test_fit_start(inference_spec, simple_model):
+    mu_init = 17
+    mu_sigma_init = 13
+
+    with simple_model:
+        if type(inference_spec()) == ADVI:
+            has_start_sigma = True
+        else:
+            has_start_sigma = False
+
+    kw = {"start": {"mu": mu_init}}
+    if has_start_sigma:
+        kw.update({"start_sigma": {"mu": mu_sigma_init}})
+
+    with simple_model:
+        inference = inference_spec(**kw)
+    trace = inference.fit(n=0).sample(10000)
+    np.testing.assert_allclose(np.mean(trace.posterior["mu"]), mu_init, rtol=0.05)
+    if has_start_sigma:
+        np.testing.assert_allclose(np.std(trace.posterior["mu"]), mu_sigma_init, rtol=0.05)
+
+
 def test_profile(inference):
     inference.run_profiling(n=100).summary()
 

pymc/variational/approximations.py

@@ -88,7 +88,7 @@ def create_shared_params(self, start=None, start_sigma=None):
     def _prepare_start_sigma(self, start_sigma):
         rho = np.zeros((self.ddim,))
         if start_sigma is not None:
-            for name, slice_, *_ in self.ordering.items():
+            for name, slice_, *_ in self.ordering.values():
                 sigma = start_sigma.get(name)
                 if sigma is not None:
                     rho[slice_] = np.log(np.exp(np.abs(sigma)) - 1.0)

pymc/variational/inference.py

@@ -257,7 +257,9 @@ def _infmean(input_array):
                     )
                 )
         else:
-            if n < 10:
+            if n == 0:
+                logger.info(f"Initialization only")
+            elif n < 10:
                 logger.info(f"Finished [100%]: Loss = {scores[-1]:,.5g}")
             else:
                 avg_loss = _infmean(scores[max(0, i - 1000) : i + 1])
@@ -466,7 +468,7 @@ class FullRankADVI(KLqp):
     random_seed: None or int
         leave None to use package global RandomStream or other
         valid value to create instance specific one
-    start: `Point`
+    start: `dict[str, np.ndarray]` or `StartDict`
         starting point for inference
 
     References
@@ -534,13 +536,11 @@ class SVGD(ImplicitGradient):
         kernel function for KSD :math:`f(histogram) -> (k(x,.), \nabla_x k(x,.))`
     temperature: float
         parameter responsible for exploration, higher temperature gives more broad posterior estimate
-    start: `dict`
+    start: `dict[str, np.ndarray]` or `StartDict`
         initial point for inference
     random_seed: None or int
         leave None to use package global RandomStream or other
         valid value to create instance specific one
-    start: `Point`
-        starting point for inference
     kwargs: other keyword arguments passed to estimator
 
     References
@@ -631,7 +631,11 @@ def __init__(self, approx=None, estimator=KSD, kernel=test_functions.rbf, **kwar
             "is often **underestimated** when using temperature = 1."
         )
         if approx is None:
-            approx = FullRank(model=kwargs.pop("model", None))
+            approx = FullRank(
+                model=kwargs.pop("model", None),
+                random_seed=kwargs.pop("random_seed", None),
+                start=kwargs.pop("start", None),
+            )
         super().__init__(estimator=estimator, approx=approx, kernel=kernel, **kwargs)
 
     def fit(