Revise basic SIR tutorial

.github/workflows/test_notebooks.yml

@@ -39,7 +39,7 @@ jobs:
         - name: Install Python packages from requirements.txt
           run: |
             pip install --upgrade pip
-            pip install -e .[test]
+            pip install -e .[test,dynamical]
 
         - name: Run Notebook Test
           run: | 

chirho/dynamical/handlers/interruption.py

@@ -126,9 +126,10 @@ def __init__(
         self,
         times: torch.Tensor,
         observation: Observation[State[T]],
+        **kwargs,
     ):
         self.observation = observation
-        super().__init__(times)
+        super().__init__(times, **kwargs)
 
     def _pyro_post_simulate(self, msg: dict) -> None:
         super()._pyro_post_simulate(msg)

chirho/dynamical/handlers/trajectory.py

@@ -19,9 +19,10 @@ class LogTrajectory(Generic[T], pyro.poutine.messenger.Messenger):
     trajectory: State[T]
     _trajectory: State[T]
 
-    def __init__(self, times: torch.Tensor):
+    def __init__(self, times: torch.Tensor, is_traced: bool = False):
         self.times = times
         self._trajectory: State[T] = State()
+        self.is_traced = is_traced
 
         # Require that the times are sorted. This is required by the index masking we do below.
         if not torch.all(self.times[1:] > self.times[:-1]):
@@ -46,6 +47,10 @@ def _pyro_post_simulate(self, msg) -> None:
         self.trajectory = self._trajectory
         self._trajectory: State[T] = State()
 
+        if self.is_traced:
+            # This adds the trajectory to the trace so that it can be accessed later.
+            [pyro.deterministic(name, value) for name, value in self.trajectory.items()]
+
     def _pyro_simulate_point(self, msg) -> None:
         # Turn a simulate that returns a state into a simulate that returns a trajectory at each of the logging_times
         dynamics, initial_state, start_time, end_time = msg["args"]