Add checkify option to console.

Author: femtomc

.flake8

@@ -2,7 +2,7 @@
 max-line-length = 88
 max-complexity = 15
 select = C,E,F,W,B,B950
-extend-ignore = E203,E501,W503,W601,E302,W605,B950,B023,F811,C901,E731
+extend-ignore = E203,E501,W503,W601,E302,W605,B950,B023,F811,C901,E731,E722
 exclude =
     .git,
     __init__.py,

src/genjax/_src/console.py

@@ -96,15 +96,40 @@ def console(
     enforce_checkify=False,
     **pretty_kwargs,
 ):
-    traceback_kwargs = {
-        "word_wrap": True,
-        "show_locals": False,
-        "max_frames": 30,
-        "suppress": [jax, plum],
-        **pretty_kwargs,
-    }
-    return GenJAXConsole(
-        Console(soft_wrap=True),
-        traceback_kwargs,
-        enforce_checkify,
-    )
+    try:
+        # Try to ignore these packages in pretty printing.
+        import asyncio
+
+        import ipykernel
+        import tornado
+        import traitlets
+
+        traceback_kwargs = {
+            "word_wrap": True,
+            "show_locals": False,
+            "max_frames": 30,
+            "suppress": [
+                jax,
+                plum,
+                asyncio,
+                tornado,
+                traitlets,
+                ipykernel,
+            ],
+            **pretty_kwargs,
+        }
+    except:
+        traceback_kwargs = {
+            "word_wrap": True,
+            "show_locals": False,
+            "max_frames": 30,
+            "suppress": [jax, plum],
+            **pretty_kwargs,
+        }
+
+    finally:
+        return GenJAXConsole(
+            Console(soft_wrap=True),
+            traceback_kwargs,
+            enforce_checkify,
+        )

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

@@ -813,7 +813,10 @@ def unmask(self):
         # contexts.
         def _check():
             check_flag = jnp.all(self.mask)
-            checkify.check(check_flag, "Mask is False, the masked value is invalid.\n")
+            checkify.check(
+                check_flag,
+                "Attempted to unmask when the mask flag is False: the masked value is invalid.\n",
+            )
 
         optional_check(_check)
         return self.value

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

@@ -205,10 +205,12 @@ def _inner(carry, xs):
 
     def _optional_out_of_bounds_check(self, count: IntArray):
         def _check():
-            check_flag = jnp.less(count + 1, self.max_length)
+            check_flag = jnp.less_equal(count + 1, self.max_length)
             checkify.check(
                 check_flag,
-                f"\nUnfoldCombinator received a length argument ({count}) longer than specified max length ({self.max_length})",
+                "UnfoldCombinator received an index argument (idx = {count}) with idx + 1 > max length ({max_length})",
+                count=jnp.array(count, copy=False),
+                max_length=jnp.array(self.max_length),
             )
 
         optional_check(_check)

src/genjax/_src/inference/translator.py

@@ -18,14 +18,17 @@
 import jax
 import jax.numpy as jnp
 import jax.tree_util as jtu
+from jax.experimental.checkify import check
 
+from genjax._src.checkify import optional_check
 from genjax._src.core.datatypes.generative import Choice
 from genjax._src.core.datatypes.generative import ChoiceMap
 from genjax._src.core.datatypes.generative import GenerativeFunction
 from genjax._src.core.datatypes.generative import Trace
 from genjax._src.core.pytree.pytree import Pytree
 from genjax._src.core.pytree.utilities import tree_grad_split
 from genjax._src.core.pytree.utilities import tree_zipper
+from genjax._src.core.typing import Bool
 from genjax._src.core.typing import Callable
 from genjax._src.core.typing import FloatArray
 from genjax._src.core.typing import PRNGKey
@@ -92,6 +95,7 @@ def safe_slogdet(v):
 class ExtendingTraceTranslator(TraceTranslator):
     choice_map_forward: Callable  # part of bijection
     choice_map_inverse: Callable  # part of bijection
+    check_bijection: Bool
     p_argdiffs: Tuple
     q_forward: GenerativeFunction
     q_forward_args: Tuple
@@ -103,7 +107,7 @@ def flatten(self):
             self.q_forward,
             self.q_forward_args,
             self.new_observations,
-        ), (self.choice_map_forward, self.choice_map_inverse)
+        ), (self.choice_map_forward, self.choice_map_inverse, self.check_bijection)
 
     @classmethod
     def new(
@@ -114,18 +118,21 @@ def new(
         new_obs: Choice,
         choice_map_forward: Callable,
         choice_map_inverse: Callable,
+        check_bijection: Bool,
     ):
         return ExtendingTraceTranslator(
             choice_map_forward,
             choice_map_inverse,
+            check_bijection,
             p_argdiffs,
             q_forward,
             q_forward_args,
             new_obs,
         )
 
     def value_and_jacobian_correction(self, forward, trace):
-        grad_tree, no_grad_tree = tree_grad_split(trace.get_choices())
+        trace_choices = trace.get_choices()
+        grad_tree, no_grad_tree = tree_grad_split(trace_choices)
 
         def _inner(differentiable):
             choices = tree_zipper(differentiable, no_grad_tree)
@@ -134,6 +141,21 @@ def _inner(differentiable):
 
         inner_jacfwd = jax.jacfwd(_inner, has_aux=True)
         J, transformed = inner_jacfwd(grad_tree)
+        if self.check_bijection:
+
+            def optional_check_bijection_is_bijection():
+                backwards = self.choice_map_inverse(transformed)
+                flattened = jtu.tree_leaves(
+                    jtu.tree_map(
+                        lambda v1, v2: jnp.all(v1 == v2),
+                        trace_choices,
+                        backwards,
+                    )
+                )
+                check_flag = jnp.all(jnp.array(flattened))
+                check(check_flag, "Bijection check failed")
+
+            optional_check(optional_check_bijection_is_bijection)
         J = stack_differentiable(J)
         (_, J_log_abs_det) = safe_slogdet(J)
         return transformed, J_log_abs_det
@@ -166,6 +188,7 @@ def extending_trace_translator(
     new_obs: ChoiceMap,
     choice_map_forward: Callable,
     choice_map_backward: Callable,
+    check_bijection=False,
 ):
     return ExtendingTraceTranslator.new(
         p_argdiffs,
@@ -174,6 +197,7 @@ def extending_trace_translator(
         new_obs,
         choice_map_forward,
         choice_map_backward,
+        check_bijection,
     )