Update pymc/step_methods/hmc/integration.py

maresb · aseyboldt · web-flow · commit 814acc5b4066 · 2022-11-13T18:38:43.000+01:00
Co-authored-by: Adrian Seyboldt &lt;aseyboldt@users.noreply.github.com&gt;
diff --git a/pymc/step_methods/hmc/integration.py b/pymc/step_methods/hmc/integration.py
@@ -26,9 +26,9 @@ class State(NamedTuple):
     q: RaveledVars
     p: RaveledVars
     v: np.ndarray
-    q_grad: Any
-    energy: Any
-    model_logp: Any
+    q_grad: np.ndarray
+    energy: float
+    model_logp: float
     index_in_trajectory: int
 
 
