Allow for StatsBijection.rmap of incomplete stat dicts

Author: michaelosthege

pymc/step_methods/compound.py

@@ -198,9 +198,12 @@ def map(self, stats_list: Sequence[Mapping[str, Any]]) -> StatsDict:
         return stats_dict
 
     def rmap(self, stats_dict: Mapping[str, Any]) -> StatsType:
-        """Split a global stats dict into a list of sampler-wise stats dicts."""
+        """Split a global stats dict into a list of sampler-wise stats dicts.
+
+        The ``stats_dict`` can be a subset of all sampler stats.
+        """
         stats_list = []
         for namemap in self._stat_groups:
-            d = {statname: stats_dict[sname] for sname, statname in namemap}
+            d = {statname: stats_dict[sname] for sname, statname in namemap if sname in stats_dict}
             stats_list.append(d)
         return stats_list

pymc/tests/step_methods/test_compound.py

@@ -130,3 +130,8 @@ def test_stats_bijection(self):
         assert isinstance(rev, list)
         assert len(rev) == len(stats_l)
         assert rev == stats_l
+        # Also rmap incomplete dicts
+        rev2 = bij.rmap({"sampler_1__a": 0})
+        assert len(rev2) == 2
+        assert len(rev2[0]) == 0
+        assert len(rev2[1]) == 1