Add test for overrides string keys

pymc-devs · Oct 13, 2021 · 4e40756 · 4e40756
1 parent 735c29e
commit 4e40756
Showing 1 changed file with 22 additions and 1 deletion.
diff --git a/pymc/tests/test_initvals.py b/pymc/tests/test_initvals.py
@@ -145,7 +145,7 @@ def test_respects_overrides(self):
             return_transformed=True,
             overrides={
                 A: at.as_tensor(2, dtype=int),
-                "B": 3,
+                B: 3,
                 C: 5,
             },
         )
@@ -154,6 +154,27 @@ def test_respects_overrides(self):
         assert np.isclose(iv["B_log__"], np.log(3))
         assert iv["C"] == 5
 
+    def test_string_overrides_work(self):
+        with pm.Model() as pmodel:
+            A = pm.Flat("A", initval=10)
+            B = pm.HalfFlat("B", initval=10)
+            C = pm.HalfFlat("C", initval=10)
+
+        fn = make_initial_point_fn(
+            model=pmodel,
+            jitter_rvs={},
+            return_transformed=True,
+            overrides={
+                "A": 1,
+                "B": 1,
+                "C_log__": 0,
+            },
+        )
+        iv = fn(0)
+        assert iv["A"] == 1
+        assert np.isclose(iv["B_log__"], 0)
+        assert iv["C_log__"] == 0
+
 
 class TestMoment:
     def test_basic(self):