fix(gpjax/thompson_sampling): add NonConjugatePosterior type annotation

JaxGaussianProcesses · Jun 27, 2024 · 0bc0435 · 0bc0435
1 parent 503d0b6
commit 0bc0435
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/gpjax/decision_making/utility_functions/thompson_sampling.py b/gpjax/decision_making/utility_functions/thompson_sampling.py
@@ -22,7 +22,10 @@
     SinglePointUtilityFunction,
 )
 from gpjax.decision_making.utils import OBJECTIVE
-from gpjax.gps import ConjugatePosterior
+from gpjax.gps import (
+    ConjugatePosterior,
+    NonConjugatePosterior,
+)
 from gpjax.typing import KeyArray
 
 
@@ -56,7 +59,7 @@ def __post_init__(self):
 
     def build_utility_function(
         self,
-        posteriors: Mapping[str, ConjugatePosterior],
+        posteriors: Mapping[str, ConjugatePosterior | NonConjugatePosterior],
         datasets: Mapping[str, Dataset],
         key: KeyArray,
     ) -> SinglePointUtilityFunction: