Merge pull request #425 from probcomp/mrb/407

Fix Map.update for #407.

Author: femtomc

Diff for: src/genjax/_src/generative_functions/combinators/vector/map_combinator.py

@@ -29,10 +29,10 @@
 from genjax._src.core.datatypes.generative import NoneSelection
 from genjax._src.core.datatypes.generative import Selection
 from genjax._src.core.datatypes.generative import Trace
-from genjax._src.core.datatypes.masks import mask
 from genjax._src.core.datatypes.tracetypes import TraceType
 from genjax._src.core.datatypes.trie import Trie
 from genjax._src.core.interpreters.staging import concrete_cond
+from genjax._src.core.transforms.incremental import tree_diff_primal
 from genjax._src.core.typing import Any
 from genjax._src.core.typing import FloatArray
 from genjax._src.core.typing import IntArray
@@ -334,30 +334,22 @@ def update(
         chm: VectorChoiceMap,
         argdiffs: Tuple,
     ):
-        def _update(key, prev, chm, argdiffs):
-            key, (retdiff, w, tr, d) = self.kernel.update(key, prev, chm, argdiffs)
-            return key, (retdiff, w, tr, d)
-
-        def _inner(key, index, prev, chm, argdiffs):
-            check = index == chm.get_index()
-            masked = mask(check, chm.inner)
-            return _update(key, prev, masked, argdiffs)
-
-        args = jtu.tree_leaves(argdiffs)
+        args = tree_diff_primal(argdiffs)
         self._static_check_broadcastable(args)
         broadcast_dim_length = self._static_broadcast_dim_length(args)
-        indices = np.array([i for i in range(0, broadcast_dim_length)])
         prev_inaxes_tree = jtu.tree_map(
             lambda v: None if v.shape == () else 0, prev.inner
         )
         key, sub_keys = slash(key, broadcast_dim_length)
+
         _, (retdiff, w, tr, discard) = jax.vmap(
-            _inner, in_axes=(0, 0, prev_inaxes_tree, 0, self.in_axes)
-        )(sub_keys, indices, prev.inner, chm, argdiffs)
+            self.kernel.update, in_axes=(0, prev_inaxes_tree, 0, self.in_axes)
+        )(sub_keys, prev.inner, chm.inner, argdiffs)
         w = jnp.sum(w)
         retval = tr.get_retval()
         scores = tr.get_score()
         map_tr = MapTrace(self, tr, args, retval, jnp.sum(scores))
+        discard = VectorChoiceMap(discard)
         return key, (retdiff, w, map_tr, discard)
 
     # The choice map passed in here is empty, but perhaps

Diff for: src/genjax/_src/inference/mcmc/metropolis_hastings.py

@@ -43,7 +43,7 @@ def apply(self, key: PRNGKey, trace: Trace, proposal_args: Tuple):
         fwd_weight = proposal_tr.get_score()
         diffs = Diff.no_change(model_args)
         key, (_, weight, new, discard) = model.update(
-            key, trace, proposal_tr.get_choices(), diffs
+            key, trace, proposal_tr.strip(), diffs
         )
         proposal_args_bwd = (new, *proposal_args)
         key, (bwd_weight, _) = self.proposal.importance(key, discard, proposal_args_bwd)

Diff for: tests/inference/mcmc/test_metropolis_hastings.py

@@ -14,6 +14,7 @@
 
 
 import jax
+import jax.numpy as jnp
 
 import genjax
 from genjax.inference.mcmc import MetropolisHastings
@@ -46,3 +47,23 @@ def proposal(nowAt, d):
                 assert tr.get_score() != new.get_score()
             else:
                 assert tr.get_score() == new.get_score()
+    
+    def test_map_combinator(self):
+        @genjax.gen
+        def model():
+            loc = genjax.normal(0., 1.) @ 'loc'
+            xs = genjax.Map(genjax.normal, in_axes=(None, 0))(loc, jnp.arange(10)) @ 'xs'
+            return xs
+
+
+        @genjax.gen
+        def proposal(choices):
+            loc = choices['loc']
+            xs = genjax.Map(genjax.normal, in_axes=(None, 0))(loc, jnp.arange(10)) @ 'xs'
+            return xs
+
+
+        key = jax.random.PRNGKey(314159)
+        key, trace = genjax.simulate(model)(key, ())
+        genjax.inference.mcmc.mh(proposal).apply(key, trace, ())
+        assert True