Merge pull request #419 from JaxGaussianProcesses/cleanup_src

Cleanup src
JaxGaussianProcesses · Nov 29, 2023 · 26aa40e · 26aa40e
2 parents e11a739 + c3aeedc
commit 26aa40e
Show file tree

Hide file tree

Showing 35 changed files with 164 additions and 1,327 deletions.
diff --git a/docs/examples/graph_kernels.py b/docs/examples/graph_kernels.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # %% [markdown]
 # # Graph Kernels
 #
@@ -23,6 +22,7 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import networkx as nx
+import optax as ox
 
 with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
@@ -154,11 +154,14 @@
 # With a posterior defined, we can now optimise the model's hyperparameters.
 
 # %%
-opt_posterior, training_history = gpx.fit_scipy(
+opt_posterior, training_history = gpx.fit(
     model=posterior,
-    objective=jit(gpx.ConjugateMLL(negative=True)),
+    objective=gpx.ConjugateMLL(negative=True),
     train_data=D,
-)
+    optim=ox.adam(learning_rate=0.01),
+    num_iters=1000,
+    key=key
+    )
 
 # %% [markdown]
 #

diff --git a/docs/examples/regression_mo.py b/docs/examples/regression_mo.py