WIP

JaxGaussianProcesses · Nov 7, 2023 · dbe8287 · dbe8287
1 parent cde11d8
commit dbe8287
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/docs/examples/intro_to_kernels.py b/docs/examples/intro_to_kernels.py
@@ -10,7 +10,6 @@
 
 config.update("jax_enable_x64", True)
 
-from jax import jit
 import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import install_import_hook, Float
@@ -235,7 +234,6 @@ def forrester(x: Float[Array, "N"]) -> Float[Array, "N"]:
 # %%
 negative_mll = gpx.objectives.ConjugateMLL(negative=True)
 negative_mll(no_opt_posterior, train_data=D)
-negative_mll = jit(negative_mll)
 
 opt_posterior, history = gpx.fit_scipy(
     model=no_opt_posterior,
@@ -536,7 +534,6 @@ def forrester(x: Float[Array, "N"]) -> Float[Array, "N"]:
 # %%
 negative_mll = gpx.objectives.ConjugateMLL(negative=True)
 negative_mll(posterior, train_data=D)
-negative_mll = jit(negative_mll)
 
 opt_posterior, history = gpx.fit(
     model=posterior,

diff --git a/gpjax/fit.py b/gpjax/fit.py
@@ -232,7 +232,9 @@ def loss(model: Module, data: Dataset) -> ScalarFloat:
         if result.success:
             print("Optimization was successful")
         else:
-            raise FailedScipyFitError("Optimization failed, try increasing max_iters or using a different optimiser.")
+            raise FailedScipyFitError(
+                "Optimization failed, try increasing max_iters or using a different optimiser."
+            )
         print(f"Final loss is {result.fun_val} after {result.num_fun_eval} iterations")
 
     # Constrained space.