Cleanup src

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]
 #