Skip trajectory building if first step diverges

Author: rlouf

aehmc/proposals.py

@@ -85,7 +85,10 @@ def progressive_uniform_sampling(
     state, energy, weight, _ = proposal
     new_state, new_energy, new_weight, _ = new_proposal
 
+    # TODO: Make the `at.isnan` check unnecessary
     p_accept = at.expit(new_weight - weight)
+    p_accept = at.where(at.isnan(p_accept), 0, p_accept)
+
     do_accept = srng.bernoulli(p_accept)
     updated_proposal = maybe_update_proposal(do_accept, proposal, new_proposal)
 

aehmc/trajectory.py

@@ -266,9 +266,21 @@ def add_one_state(
             state[1],
             0,
         )
+        full_initial_state = (
+            *proposal[0],
+            proposal[1],
+            proposal[2],
+            proposal[3],
+            *state,
+            momentum_sum,
+            *termination_state,
+            at.as_tensor(1, dtype=np.int32),
+            is_diverging,
+            np.array(False),
+        )
 
         steps = at.arange(1, 1 + max_num_steps)
-        traj, updates = aesara.scan(
+        trajectory, updates = aesara.scan(
             add_one_state,
             outputs_info=(
                 *proposal[0],
@@ -284,19 +296,28 @@ def add_one_state(
             ),
             sequences=steps,
         )
+        full_last_state = tuple([state[-1] for state in trajectory])
+
+        # We build the trajectory iff the first step is not diverging
+        full_state = ifelse(is_diverging, full_initial_state, full_last_state)
 
         new_proposal = (
-            (traj[0][-1], traj[1][-1], traj[2][-1], traj[3][-1]),
-            traj[4][-1],
-            traj[5][-1],
-            traj[6][-1],
+            (full_state[0], full_state[1], full_state[2], full_state[3]),
+            full_state[4],
+            full_state[5],
+            full_state[6],
+        )
+        new_state = (full_state[7], full_state[8], full_state[9], full_state[10])
+        subtree_momentum_sum = full_state[11]
+        new_termination_state = (
+            full_state[12],
+            full_state[13],
+            full_state[14],
+            full_state[15],
         )
-        new_state = (traj[7][-1], traj[8][-1], traj[9][-1], traj[10][-1])
-        subtree_momentum_sum = traj[11][-1]
-        new_termination_state = (traj[12][-1], traj[13][-1], traj[14][-1], traj[15][-1])
-        trajectory_length = traj[-3][-1]  # defined as the number of steps taken
-        is_diverging = traj[-2][-1] | is_diverging
-        has_terminated = traj[-1][-1]
+        trajectory_length = full_state[-3]
+        is_diverging = full_state[-2]
+        has_terminated = full_state[-1]
 
         return (
             new_proposal,

tests/test_trajectory.py

@@ -71,7 +71,8 @@ def test_static_integration(example):
     "case",
     [
         (0.0000001, False, False),
-        (1000, True, True),
+        (1000, True, False),
+        (1e100, True, False),
     ],
 )
 def test_dynamic_integration(case):
@@ -136,7 +137,7 @@ def potential_fn(x):
 @pytest.mark.parametrize(
     "step_size, should_diverge, should_turn, expected_doublings",
     [
-        (100000.0, True, True, 1),
+        (100000.0, True, False, 1),
         (0.0000001, False, False, 10),
         (1.0, False, True, 1),
     ],