Local.

Author: femtomc

.gitattributes

@@ -5,3 +5,4 @@ site/* linguist-documentation
 *.html linguist-vendored
 *.js linguist-vendored
 *.ipynb linguist-vendored
+*.css linguist-vendored

mkdocs.yml

@@ -18,11 +18,11 @@ repo_url: https://github.com/probcomp/genjax
 copyright: Copyright © 2023 MIT Probabilistic Computing Project
 
 nav:
-  - Home: index.md
+  - Home: 
+      - index.md
+      - Gen's concepts: genjax/concepts/generative_functions.md
+      - Diff against Gen.jl: genjax/diff_jl.md
   - Inference notebooks: genjax/notebooks.md
-  - Concepts:
-      - Generative functions: genjax/concepts/generative_functions.md
-      - Diff with Gen.jl: genjax/concepts/diff_jl.md
   - Language apéritifs: genjax/language_aperitifs.md
   - Public API reference:
       - Core: genjax/library/core.md

src/genjax/_src/core/datatypes/trie.py

@@ -289,6 +289,22 @@ def get_subtrees_shallow(self):
         return self.trie.get_subtrees_shallow()
 
 
+###################
+# TrieConvertable #
+###################
+
+# A mixin: denotes that a choice map can be converted to a TrieChoiceMap
+
+
+@dataclass
+class TrieConvertable:
+    def convert(self) -> TrieChoiceMap:
+        new = TrieChoiceMap.new()
+        for (k, v) in self.get_submaps_shallow():
+            pass
+        return new
+
+
 ##############
 # Shorthands #
 ##############

src/genjax/_src/generative_functions/combinators/vector/map_combinator.py

@@ -28,9 +28,11 @@
 from genjax._src.core.datatypes.generative import GenerativeFunction
 from genjax._src.core.datatypes.masks import mask
 from genjax._src.core.datatypes.tracetypes import TraceType
+from genjax._src.core.datatypes.trie import TrieChoiceMap
 from genjax._src.core.interpreters.staging import concrete_cond
 from genjax._src.core.typing import Any
 from genjax._src.core.typing import FloatArray
+from genjax._src.core.typing import Int
 from genjax._src.core.typing import IntArray
 from genjax._src.core.typing import PRNGKey
 from genjax._src.core.typing import Tuple
@@ -194,6 +196,14 @@ def _padder(self, v, key_len):
             np.pad(v, pad_axes) if isinstance(v, np.ndarray) else jnp.pad(v, pad_axes)
         )
 
+    def _static_check_trie_index_compatible(
+        self, chm: TrieChoiceMap, broadcast_dim_length: Int
+    ):
+        for (k, _) in chm.get_subtrees_shallow():
+            assert isinstance(k, int)
+            # TODO: pull outside loop, just check the last address.
+            assert k < broadcast_dim_length
+
     def _importance_vcm(self, key, chm, args):
         def _importance(key, chm, args):
             return self.kernel.importance(key, chm, args)
@@ -218,6 +228,16 @@ def _inner(key, index, chm, args):
         map_tr = VectorTrace(self, indices, tr, args, retval, scores)
         return key, (w, map_tr)
 
+    # Implements a conversion from `TrieChoiceMap`.
+    def _importance_tchm(self, key, chm, args):
+        broadcast_dim_length = self._static_broadcast_dim_length(args)
+        self._static_check_trie_index_compatible(chm, broadcast_dim_length)
+
+        # Okay, so the TrieChoiceMap has an address hierarchy which is compatible with the index structure of the MapCombinator choices.
+        # Let's coerce TrieChoiceMap into VectorChoiceMap and then just call `_importance_vcm`.
+        vector_chm = self._coerce_to_vector_chm(chm)
+        return self._importance_vcm(key, vector_chm, args)
+
     def _importance_empty(self, key, _, args):
         key, map_tr = self.simulate(key, args)
         w = 0.0
@@ -227,7 +247,7 @@ def _importance_empty(self, key, _, args):
     def importance(
         self,
         key: PRNGKey,
-        chm: Union[EmptyChoiceMap, VectorChoiceMap],
+        chm: Union[EmptyChoiceMap, TrieChoiceMap, VectorChoiceMap],
         args: Tuple,
         **_,
     ) -> Tuple[PRNGKey, Tuple[FloatArray, VectorTrace]]:
@@ -237,6 +257,8 @@ def importance(
         # Note: these branches are resolved at tracing time.
         if isinstance(chm, VectorChoiceMap):
             return self._importance_vcm(key, chm, args)
+        elif isinstance(chm, TrieChoiceMap):
+            return self._importance_tchm(key, chm, args)
         else:
             assert isinstance(chm, EmptyChoiceMap)
             return self._importance_empty(key, chm, args)