Refactor docstrings

benchmarks/objectives.py

@@ -22,7 +22,7 @@ def setup(self, n_datapoints: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Gaussian(num_datapoints=self.data.n)
         self.objective = gpx.ConjugateMLL()
         self.posterior = self.prior * self.likelihood
@@ -48,7 +48,7 @@ def setup(self, n_datapoints: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Bernoulli(num_datapoints=self.data.n)
         self.objective = gpx.LogPosteriorDensity()
         self.posterior = self.prior * self.likelihood
@@ -75,7 +75,7 @@ def setup(self, n_datapoints: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Poisson(num_datapoints=self.data.n)
         self.objective = gpx.LogPosteriorDensity()
         self.posterior = self.prior * self.likelihood

benchmarks/predictions.py

@@ -21,7 +21,7 @@ def setup(self, n_test: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Gaussian(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
         key, subkey = jr.split(key)
@@ -46,7 +46,7 @@ def setup(self, n_test: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Bernoulli(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
         key, subkey = jr.split(key)
@@ -71,7 +71,7 @@ def setup(self, n_test: int, n_dims: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(n_dims)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Bernoulli(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
         key, subkey = jr.split(key)

benchmarks/sparse.py

@@ -19,7 +19,7 @@ def setup(self, n_datapoints: int, n_inducing: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(1)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Gaussian(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
 

benchmarks/stochastic.py

@@ -20,7 +20,7 @@ def setup(self, n_datapoints: int, n_inducing: int, batch_size: int):
         self.data = gpx.Dataset(X=self.X, y=self.y)
         kernel = gpx.kernels.RBF(active_dims=list(range(1)))
         meanf = gpx.mean_functions.Constant()
-        self.prior = gpx.Prior(kernel=kernel, mean_function=meanf)
+        self.prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf)
         self.likelihood = gpx.likelihoods.Gaussian(num_datapoints=self.data.n)
         self.posterior = self.prior * self.likelihood
 

gpjax/fit.py

@@ -70,17 +70,17 @@ def fit(  # noqa: PLR0913
         >>> D = gpx.Dataset(X, y)
         >>>
         >>> # (2) Define your model:
-        >>> class LinearModel(gpx.Module):
-                weight: float = gpx.param_field()
-                bias: float = gpx.param_field()
+        >>> class LinearModel(gpx.base.Module):
+                weight: float = gpx.base.param_field()
+                bias: float = gpx.base.param_field()
 
                 def __call__(self, x):
                     return self.weight * x + self.bias
 
         >>> model = LinearModel(weight=1.0, bias=1.0)
         >>>
         >>> # (3) Define your loss function:
-        >>> class MeanSquareError(gpx.AbstractObjective):
+        >>> class MeanSquareError(gpx.objectives.AbstractObjective):
                 def evaluate(self, model: LinearModel, train_data: gpx.Dataset) -> float:
                     return jnp.mean((train_data.y - model(train_data.X)) ** 2)
         >>>

gpjax/gps.py

@@ -149,7 +149,7 @@ class Prior(AbstractPrior):
 
         >>> kernel = gpx.kernels.RBF()
         >>> meanf = gpx.mean_functions.Zero()
-        >>> prior = gpx.Prior(mean_function=meanf, kernel = kernel)
+        >>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
     ```
     """
 
@@ -183,7 +183,7 @@ def __mul__(self, other):
             >>>
             >>> meanf = gpx.mean_functions.Zero()
             >>> kernel = gpx.kernels.RBF()
-            >>> prior = gpx.Prior(mean_function=meanf, kernel = kernel)
+            >>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
             >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=100)
             >>>
             >>> prior * likelihood
@@ -244,7 +244,7 @@ def predict(self, test_inputs: Num[Array, "N D"]) -> ReshapedGaussianDistributio
             >>>
             >>> kernel = gpx.kernels.RBF()
             >>> meanf = gpx.mean_functions.Zero()
-            >>> prior = gpx.Prior(mean_function=meanf, kernel = kernel)
+            >>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
             >>>
             >>> prior.predict(jnp.linspace(0, 1, 100))
         ```
@@ -310,7 +310,7 @@ def sample_approx(
             >>>
             >>> meanf = gpx.mean_functions.Zero()
             >>> kernel = gpx.kernels.RBF()
-            >>> prior = gpx.Prior(mean_function=meanf, kernel = kernel)
+            >>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
             >>>
             >>> sample_fn = prior.sample_approx(10, key)
             >>> sample_fn(jnp.linspace(0, 1, 100).reshape(-1, 1))
@@ -434,7 +434,7 @@ class ConjugatePosterior(AbstractPosterior):
             >>> import gpjax as gpx
             >>> import jax.numpy as jnp
 
-            >>> prior = gpx.Prior(
+            >>> prior = gpx.gps.Prior(
                     mean_function = gpx.mean_functions.Zero(),
                     kernel = gpx.kernels.RBF()
                 )
@@ -482,8 +482,8 @@ def predict(
                 >>> D = gpx.Dataset(X=xtrain, y=ytrain)
                 >>> xtest = jnp.linspace(0, 1).reshape(-1, 1)
                 >>>
-                >>> prior = gpx.Prior(mean_function = gpx.Zero(), kernel = gpx.RBF())
-                >>> posterior = prior * gpx.Gaussian(num_datapoints = D.n)
+                >>> prior = gpx.gps.Prior(mean_function = gpx.mean_functions.Zero(), kernel = gpx.kernels.RBF())
+                >>> posterior = prior * gpx.likelihoods.Gaussian(num_datapoints = D.n)
                 >>> predictive_dist = posterior(xtest, D)
             ```
 

gpjax/objectives.py

@@ -100,10 +100,10 @@ def step(
             >>> meanf = gpx.mean_functions.Constant()
             >>> kernel = gpx.kernels.RBF()
             >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n)
-            >>> prior = gpx.Prior(mean_function = meanf, kernel=kernel)
+            >>> prior = gpx.gps.Prior(mean_function = meanf, kernel=kernel)
             >>> posterior = prior * likelihood
             >>>
-            >>> mll = gpx.ConjugateMLL(negative=True)
+            >>> mll = gpx.objectives.ConjugateMLL(negative=True)
             >>> mll(posterior, train_data = D)
         ```
 
@@ -112,13 +112,13 @@ def step(
         marginal log-likelihood. This can be realised through
 
         ```python
-            mll = gpx.ConjugateMLL(negative=True)
+            mll = gpx.objectives.ConjugateMLL(negative=True)
         ```
 
         For optimal performance, the marginal log-likelihood should be ``jax.jit``
         compiled.
         ```python
-            mll = jit(gpx.ConjugateMLL(negative=True))
+            mll = jit(gpx.objectives.ConjugateMLL(negative=True))
         ```
 
         Args:

tests/test_objectives.py

@@ -64,7 +64,8 @@ def test_conjugate_mll(
 
     # Build model
     p = Prior(
-        kernel=gpx.RBF(active_dims=list(range(num_dims))), mean_function=gpx.Constant()
+        kernel=gpx.kernels.RBF(active_dims=list(range(num_dims))),
+        mean_function=gpx.mean_functions.Constant(),
     )
     likelihood = Gaussian(num_datapoints=num_datapoints)
     post = p * likelihood
@@ -93,7 +94,8 @@ def test_non_conjugate_mll(
 
     # Build model
     p = Prior(
-        kernel=gpx.RBF(active_dims=list(range(num_dims))), mean_function=gpx.Constant()
+        kernel=gpx.kernels.RBF(active_dims=list(range(num_dims))),
+        mean_function=gpx.mean_functions.Constant(),
     )
     likelihood = Bernoulli(num_datapoints=num_datapoints)
     post = p * likelihood
@@ -129,7 +131,8 @@ def test_collapsed_elbo(
     )
 
     p = Prior(
-        kernel=gpx.RBF(active_dims=list(range(num_dims))), mean_function=gpx.Constant()
+        kernel=gpx.kernels.RBF(active_dims=list(range(num_dims))),
+        mean_function=gpx.mean_functions.Constant(),
     )
     likelihood = Gaussian(num_datapoints=num_datapoints)
     q = gpx.CollapsedVariationalGaussian(posterior=p * likelihood, inducing_inputs=z)
@@ -169,7 +172,8 @@ def test_elbo(
     )
 
     p = Prior(
-        kernel=gpx.RBF(active_dims=list(range(num_dims))), mean_function=gpx.Constant()
+        kernel=gpx.kernels.RBF(active_dims=list(range(num_dims))),
+        mean_function=gpx.mean_functions.Constant(),
     )
     if binary:
         likelihood = Bernoulli(num_datapoints=num_datapoints)