Allow to pass trainable inducing inputs to AbstractVariationalGaussian

JaxGaussianProcesses · Nov 1, 2024 · 3152cde · 3152cde
1 parent 8f14cbd
commit 3152cde
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/gpjax/variational_families.py b/gpjax/variational_families.py
@@ -108,10 +108,17 @@ class AbstractVariationalGaussian(AbstractVariationalFamily[L]):
     def __init__(
         self,
         posterior: AbstractPosterior[P, L],
-        inducing_inputs: Float[Array, "N D"],
+        inducing_inputs: tp.Union[
+            Float[Array, "N D"],
+            Real,
+            Static,
+        ],
         jitter: ScalarFloat = 1e-6,
     ):
-        self.inducing_inputs = Static(inducing_inputs)
+        if not isinstance(inducing_inputs, (Real, Static)):
+            inducing_inputs = Real(inducing_inputs)
+
+        self.inducing_inputs = inducing_inputs
         self.jitter = jitter
 
         super().__init__(posterior)