re format

doc/sphinx/notebooks/advanced/fast_predictions.pct.py

@@ -78,12 +78,13 @@
 
 # +
 import numpy as np
-
-import gpflow
-from gpflow.ci_utils import reduce_in_tests
 import tensorflow as tf
 import tensorflow_probability as tfp
 from tensorflow_probability import distributions as tfd
+
+import gpflow
+from gpflow.ci_utils import reduce_in_tests
+
 f64 = gpflow.utilities.to_default_float
 
 
@@ -127,9 +128,6 @@
 model.predict_y_faster(Xnew, posteriors=posterior)
 
 
-
-
-
 # ## SVGP Example
 #
 # Likewise, we will construct an SVGP model to demonstrate the faster predictions from using the cached data in the GPFlow posterior classes.
@@ -196,15 +194,14 @@
 model.predict_y_faster(Xnew, posteriors=posterior)
 
 
-
 # ## MCMC for GPR Example
 #
 # Faster predictions for the sampled hyperparameters in Gaussian process regression
 # GPR with HMC
 
 # Model setup and Hyperparameters sampling. See the [MCMC (Markov Chain Monte Carlo)](../advanced/mcmc.ipynb) for more details.
 # %%
-data = (X, Y) 
+data = (X, Y)
 kernel = gpflow.kernels.Matern52(lengthscales=0.3)
 mean_function = gpflow.mean_functions.Linear(1.0, 0.0)
 model = gpflow.models.GPR(data, kernel, mean_function, noise_variance=0.01)
@@ -219,7 +216,7 @@
 gpflow.utilities.print_summary(model)
 
 # Sampling hyperparameters
-num_burnin_steps = 5 
+num_burnin_steps = 5
 num_samples = 100
 
 # Note that here we need model.trainable_parameters, not trainable_variables - only parameters can have priors!
@@ -228,11 +225,18 @@
 )
 
 hmc = tfp.mcmc.HamiltonianMonteCarlo(
-    target_log_prob_fn=hmc_helper.target_log_prob_fn, num_leapfrog_steps=10, step_size=0.01
+    target_log_prob_fn=hmc_helper.target_log_prob_fn,
+    num_leapfrog_steps=10,
+    step_size=0.01,
 )
 adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
-    hmc, num_adaptation_steps=10, target_accept_prob=f64(0.75), adaptation_rate=0.1
+    hmc,
+    num_adaptation_steps=10,
+    target_accept_prob=f64(0.75),
+    adaptation_rate=0.1,
 )
+
+
 @tf.function
 def run_chain_fn():
     return tfp.mcmc.sample_chain(
@@ -243,6 +247,7 @@ def run_chain_fn():
         trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
     )
 
+
 samples, traces = run_chain_fn()
 parameter_samples = hmc_helper.convert_to_constrained_values(samples)
 
@@ -271,10 +276,10 @@ def run_chain_fn():
         var.assign(var_samples[i])
     model.predict_f_loaded_cache(Xnew, Cache[i])
 
-# Now, the faster predictions can be made as long as the variables Cache and samples are stored in the local. 
+# Now, the faster predictions can be made as long as the variables Cache and samples are stored in the local.
 # %%
 # For example, if we delete the old model
-del(model)
+del model
 # build a new model called model2 with identical initial setup as the old model
 kernel = gpflow.kernels.Matern52(lengthscales=0.3)
 mean_function = gpflow.mean_functions.Linear(1.0, 0.0)
@@ -289,7 +294,7 @@ def run_chain_fn():
 model2.mean_function.b.prior = tfd.Normal(f64(0.0), f64(10.0))
 
 # %%timeit
-# we can still make faster predictions through variables Cache and samples 
+# we can still make faster predictions through variables Cache and samples
 for _ in range(0, 100):
     i = np.random.randint(0, 100)
     for var, var_samples in zip(hmc_helper.current_state, samples):
@@ -305,7 +310,7 @@ def run_chain_fn():
     model2.predict_y_loaded_cache(Xnew, Cache[i])
 
 # ## MCMC for Gaussian process models (GPMC)
-# This is similar as the case of GPR. 
+# This is similar as the case of GPR.
 
 # %%
 # Create some data
@@ -328,9 +333,11 @@ def run_chain_fn():
 
 # sampling hyperparameters
 optimizer = gpflow.optimizers.Scipy()
-maxiter =300 
+maxiter = 300
 _ = optimizer.minimize(
-    model.training_loss, model.trainable_variables, options=dict(maxiter=maxiter)
+    model.training_loss,
+    model.trainable_variables,
+    options=dict(maxiter=maxiter),
 )
 num_burnin_steps = 10
 num_samples = 100
@@ -341,13 +348,19 @@ def run_chain_fn():
 )
 
 hmc = tfp.mcmc.HamiltonianMonteCarlo(
-    target_log_prob_fn=hmc_helper.target_log_prob_fn, num_leapfrog_steps=10, step_size=0.01
+    target_log_prob_fn=hmc_helper.target_log_prob_fn,
+    num_leapfrog_steps=10,
+    step_size=0.01,
 )
 
 adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
-    hmc, num_adaptation_steps=10, target_accept_prob=f64(0.75), adaptation_rate=0.1
+    hmc,
+    num_adaptation_steps=10,
+    target_accept_prob=f64(0.75),
+    adaptation_rate=0.1,
 )
 
+
 @tf.function
 def run_chain_fn():
     return tfp.mcmc.sample_chain(
@@ -358,6 +371,7 @@ def run_chain_fn():
         trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
     )
 
+
 samples, _ = run_chain_fn()
 
 # %%timeit
@@ -381,6 +395,3 @@ def run_chain_fn():
     for var, var_samples in zip(hmc_helper.current_state, samples):
         var.assign(var_samples[i])
     model.predict_f_loaded_cache(Xtest, Cache[i])
-
-
-

gpflow/conditionals/conditionals.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import tensorflow as tf
 from check_shapes import check_shapes
@@ -98,14 +98,13 @@ def _sparse_conditional(
     "return[1]: [batch..., N, R, R] if (not full_cov) and full_output_cov",
     "return[1]: [batch..., N, R, N, R] if full_cov and full_output_cov",
 )
-
 def _dense_conditional(
     Xnew: tf.Tensor,
     X: tf.Tensor,
     kernel: Kernel,
     f: tf.Tensor,
     *,
-    Cache: tf.Tensor = None,
+    Cache: Optional[Tuple[tf.Tensor, ...]] = None,
     full_cov: bool = False,
     full_output_cov: bool = False,
     q_sqrt: Optional[tf.Tensor] = None,
@@ -161,4 +160,4 @@ def _dense_conditional(
         posterior.cache = Cache
         return posterior.predict_f(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)
     else:
-        return posterior.fused_predict_f(Xnew, full_cov=full_cov, full_output_cov=full_output_cov) 
+        return posterior.fused_predict_f(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)

gpflow/models/gpmc.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import numpy as np
 import tensorflow as tf
 import tensorflow_probability as tfp
 from check_shapes import check_shapes, inherit_check_shapes
 
+from .. import posteriors
 from ..base import InputData, MeanAndVariance, Parameter, RegressionData
 from ..conditionals import conditional
 from ..config import default_float, default_jitter
@@ -29,7 +30,7 @@
 from .model import GPModel
 from .training_mixins import InternalDataTrainingLossMixin
 from .util import data_input_to_tensor
-from .. import posteriors
+
 
 class GPMC_deprecated(GPModel, InternalDataTrainingLossMixin):
     @check_shapes(
@@ -128,8 +129,6 @@ def predict_f(
         return mu + self.mean_function(Xnew), var
 
 
-
-
 class GPMC_with_posterior(GPMC_deprecated):
     """
     This is an implementation of GPMC that provides a posterior() method that
@@ -164,10 +163,10 @@ def posterior(
             kernel=self.kernel,
             X=X_data,
             q_mu=self.V,
-            q_sqrt=None, 
+            q_sqrt=None,
             white=True,
-            precompute_cache=precompute_cache
-        ) 
+            precompute_cache=precompute_cache,
+        )
 
     @check_shapes(
         "Xnew: [batch..., N, D]",
@@ -177,9 +176,13 @@ def posterior(
         "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
         "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
         "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
-    ) 
+    )
     def predict_f_loaded_cache(
-        self, Xnew: InputData, Cache: InputData,full_cov: bool = False, full_output_cov: bool = False
+        self,
+        Xnew: InputData,
+        Cache: Optional[Tuple[tf.Tensor, ...]],
+        full_cov: bool = False,
+        full_output_cov: bool = False,
     ) -> MeanAndVariance:
         """
         Xnew is a data matrix, point at which we want to predict
@@ -195,11 +198,17 @@ def predict_f_loaded_cache(
 
         X_data, _Y_data = self.data
         mu, var = conditional(
-            Xnew, X_data, self.kernel, self.V, Cache = Cache, full_cov=full_cov, q_sqrt=None, white=True
+            Xnew,
+            X_data,
+            self.kernel,
+            self.V,
+            Cache=Cache,
+            full_cov=full_cov,
+            q_sqrt=None,
+            white=True,
         )
         return mu + self.mean_function(Xnew), var
 
-
     @check_shapes(
         "Xnew: [batch..., N, D]",
         "Cache[0]: [M, M]",
@@ -209,9 +218,12 @@ def predict_f_loaded_cache(
         "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
         "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
     )
-
     def predict_y_loaded_cache(
-        self, Xnew: InputData, Cache: InputData, full_cov: bool = False, full_output_cov: bool = False
+        self,
+        Xnew: InputData,
+        Cache: Optional[Tuple[tf.Tensor, ...]],
+        full_cov: bool = False,
+        full_output_cov: bool = False,
     ) -> MeanAndVariance:
         """
         For backwards compatibility, GPR's predict_f uses the fused (no-cache)
@@ -221,12 +233,14 @@ def predict_y_loaded_cache(
             model.posterior().predict_f(Xnew, ...)
         """
 
-        f_mean, f_var = self.predict_f_loaded_cache(Xnew, Cache, full_cov=full_cov, full_output_cov=full_output_cov)
+        f_mean, f_var = self.predict_f_loaded_cache(
+            Xnew, Cache, full_cov=full_cov, full_output_cov=full_output_cov
+        )
 
         return self.likelihood.predict_mean_and_var(Xnew, f_mean, f_var)
 
-    
+
 class GPMC(GPMC_with_posterior):
     # subclassed to ensure __class__ == "GPMC"
 
-    __doc__ = GPMC_deprecated.__doc__  # Use documentation from GPMC_deprecated.
+    __doc__ = GPMC_deprecated.__doc__  # Use documentation from GPMC_deprecated.