Use the correct givens in Model.set_initval

Author: brandonwillard

pymc3/model.py

@@ -938,8 +938,7 @@ def set_initval(self, rv_var, initval):
 
         if initval is None or transform:
             # Sample/evaluate this using the existing initial values, and
-            # with the least amount of affect on the RNGs involved (i.e. no
-            # in-placing)
+            # with the least effect on the RNGs involved (i.e. no in-placing)
             from aesara.compile.mode import Mode, get_mode
 
             mode = get_mode(None)
@@ -950,8 +949,21 @@ def set_initval(self, rv_var, initval):
                 value = initval if initval is not None else rv_var
                 rv_var = transform.forward(rv_var, value)
 
+            def initval_to_rvval(value_var, value):
+                rv_var = self.values_to_rvs[value_var]
+                initval = value_var.type.make_constant(value)
+                transform = getattr(value_var.tag, "transform", None)
+                if transform:
+                    return transform.backward(rv_var, initval)
+                else:
+                    return initval
+
+            givens = {
+                self.values_to_rvs[k]: initval_to_rvval(k, v)
+                for k, v in self.initial_values.items()
+            }
             initval_fn = aesara.function(
-                [], rv_var, mode=mode, givens=self.initial_values, on_unused_input="ignore"
+                [], rv_var, mode=mode, givens=givens, on_unused_input="ignore"
             )
             initval = initval_fn()
 

pymc3/tests/test_model.py

@@ -632,3 +632,27 @@ def test_invalid_variable_name(self):
         model.update_start_vals(start, model.initial_point)
         with pytest.raises(KeyError):
             model.check_start_vals(start)
+
+
+def test_set_initval():
+    # Make sure the dependencies between variables are maintained when
+    # generating initial values
+    rng = np.random.RandomState(392)
+
+    with pm.Model(rng_seeder=rng) as model:
+        eta = pm.Uniform("eta", 1.0, 2.0, size=(1, 1))
+        mu = pm.Normal("mu", sd=eta, initval=[[100]])
+        alpha = pm.HalfNormal("alpha", initval=100)
+        value = pm.NegativeBinomial("value", mu=mu, alpha=alpha)
+
+    assert np.array_equal(model.initial_values[model.rvs_to_values[mu]], np.array([[100.0]]))
+    np.testing.assert_almost_equal(model.initial_values[model.rvs_to_values[alpha]], np.log(100))
+    assert 50 < model.initial_values[model.rvs_to_values[value]] < 150
+
+    # `Flat` cannot be sampled, so let's make sure that doesn't break initial
+    # value computations
+    with pm.Model() as model:
+        x = pm.Flat("x")
+        y = pm.Normal("y", x, 1)
+
+    assert model.rvs_to_values[y] in model.initial_values