gpjax 0.13.1__py3-none-any.whl → 0.13.3__py3-none-any.whl

gpjax/__init__.py

@@ -38,9 +38,9 @@ from gpjax.fit import (
 
 __license__ = "MIT"
 __description__ = "Gaussian processes in JAX and Flax"
-__url__ = "https://github.com/JaxGaussianProcesses/GPJax"
-__contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.13.1"
+__url__ = "https://github.com/thomaspinder/GPJax"
+__contributors__ = "https://github.com/thomaspinder/GPJax/graphs/contributors"
+__version__ = "0.13.3"
 
 __all__ = [
     "gps",

gpjax/dataset.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/fit.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2023 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/gps.py

@@ -4,7 +4,7 @@
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+# http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
@@ -13,10 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 # from __future__ import annotations
+
 from abc import abstractmethod
+from typing import Literal
 
 import beartype.typing as tp
 from flax import nnx
+import jax
 import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import (
@@ -35,10 +38,13 @@ from gpjax.likelihoods import (
 )
 from gpjax.linalg import (
     Dense,
+    Diagonal,
     psd,
     solve,
 )
-from gpjax.linalg.operations import lower_cholesky
+from gpjax.linalg.operations import (
+    lower_cholesky,
+)
 from gpjax.linalg.utils import add_jitter
 from gpjax.mean_functions import AbstractMeanFunction
 from gpjax.parameters import (
@@ -77,7 +83,12 @@ class AbstractPrior(nnx.Module, tp.Generic[M, K]):
         self.mean_function = mean_function
         self.jitter = jitter
 
-    def __call__(self, test_inputs: Num[Array, "N D"]) -> GaussianDistribution:
+    def __call__(
+        self,
+        test_inputs: Num[Array, "N D"],
+        *,
+        return_covariance_type: Literal["dense", "diagonal"] = "dense",
+    ) -> GaussianDistribution:
         r"""Evaluate the Gaussian process at the given points.
 
         The output of this function is a
@@ -91,15 +102,27 @@ class AbstractPrior(nnx.Module, tp.Generic[M, K]):
 
         Args:
             test_inputs: Input locations where the GP should be evaluated.
+            return_covariance_type: Literal denoting whether to return the full covariance
+                of the joint predictive distribution at the test_inputs (dense)
+                or just the the standard-deviation of the predictive distribution at
+                the test_inputs.
 
         Returns:
             GaussianDistribution: A multivariate normal random variable representation
                 of the Gaussian process.
         """
-        return self.predict(test_inputs)
+        return self.predict(
+            test_inputs,
+            return_covariance_type=return_covariance_type,
+        )
 
     @abstractmethod
-    def predict(self, test_inputs: Num[Array, "N D"]) -> GaussianDistribution:
+    def predict(
+        self,
+        test_inputs: Num[Array, "N D"],
+        *,
+        return_covariance_type: Literal["dense", "diagonal"] = "dense",
+    ) -> GaussianDistribution:
         r"""Evaluate the predictive distribution.
 
         Compute the latent function's multivariate normal distribution for a
@@ -108,6 +131,10 @@ class AbstractPrior(nnx.Module, tp.Generic[M, K]):
 
         Args:
             test_inputs: Input locations where the GP should be evaluated.
+            return_covariance_type: Literal denoting whether to return the full covariance
+                of the joint predictive distribution at the test_inputs (dense)
+                or just the the standard-deviation of the predictive distribution at
+                the test_inputs.
 
         Returns:
             GaussianDistribution: A multivariate normal random variable representation
@@ -123,8 +150,8 @@ class Prior(AbstractPrior[M, K]):
     r"""A Gaussian process prior object.
 
     The GP is parameterised by a
-    [mean](https://docs.jaxgaussianprocesses.com/api/mean_functions/)
-    and [kernel](https://docs.jaxgaussianprocesses.com/api/kernels/base/)
+    [mean](https://docs.thomaspinder.com/api/mean_functions/)
+    and [kernel](https://docs.thomaspinder.com/api/kernels/base/)
     function.
 
     A Gaussian process prior parameterised by a mean function $m(\cdot)$ and a kernel
@@ -220,7 +247,12 @@ class Prior(AbstractPrior[M, K]):
         """
         return self.__mul__(other)
 
-    def predict(self, test_inputs: Num[Array, "N D"]) -> GaussianDistribution:
+    def predict(
+        self,
+        test_inputs: Num[Array, "N D"],
+        *,
+        return_covariance_type: Literal["dense", "diagonal"] = "dense",
+    ) -> GaussianDistribution:
         r"""Compute the predictive prior distribution for a given set of
         parameters. The output of this function is a function that computes
         a TFP distribution for a given set of inputs.
@@ -241,17 +273,43 @@ class Prior(AbstractPrior[M, K]):
         Args:
             test_inputs (Float[Array, "N D"]): The inputs at which to evaluate the
                 prior distribution.
+            return_covariance_type: Literal denoting whether to return the full covariance
+                of the joint predictive distribution at the test_inputs (dense)
+                or just the the standard-deviation of the predictive distribution at
+                the test_inputs.
 
         Returns:
             GaussianDistribution: A multivariate normal random variable representation
                 of the Gaussian process.
         """
+
+        def _return_full_covariance(
+            t: Num[Array, "N D"],
+        ) -> Dense:
+            Kxx = self.kernel.gram(t)
+            Kxx_dense = add_jitter(Kxx.to_dense(), self.jitter)
+            Kxx = psd(Dense(Kxx_dense))
+            return Kxx
+
+        def _return_diagonal_covariance(
+            t: Num[Array, "N D"],
+        ) -> Dense:
+            Kxx = self.kernel.diagonal(t).diagonal
+            Kxx += self.jitter
+            Kxx = psd(Dense(Diagonal(Kxx).to_dense()))
+            return Kxx
+
         mean_at_test = self.mean_function(test_inputs)
-        Kxx = self.kernel.gram(test_inputs)
-        Kxx_dense = add_jitter(Kxx.to_dense(), self.jitter)
-        Kxx = psd(Dense(Kxx_dense))
+        cov = jax.lax.cond(
+            return_covariance_type == "dense",
+            _return_full_covariance,
+            _return_diagonal_covariance,
+            test_inputs,
+        )
 
-        return GaussianDistribution(jnp.atleast_1d(mean_at_test.squeeze()), Kxx)
+        return GaussianDistribution(
+            loc=jnp.atleast_1d(mean_at_test.squeeze()), scale=cov
+        )
 
     def sample_approx(
         self,
@@ -329,7 +387,7 @@ P = tp.TypeVar("P", bound=AbstractPrior)
 
 #######################
 # GP Posteriors
-#######################
+#######################from gpjax.linalg.operators import LinearOperator
 class AbstractPosterior(nnx.Module, tp.Generic[P, L]):
     r"""Abstract Gaussian process posterior.
 
@@ -356,7 +414,11 @@ class AbstractPosterior(nnx.Module, tp.Generic[P, L]):
         self.jitter = jitter
 
     def __call__(
-        self, test_inputs: Num[Array, "N D"], train_data: Dataset
+        self,
+        test_inputs: Num[Array, "N D"],
+        train_data: Dataset,
+        *,
+        return_covariance_type: Literal["dense", "diagonal"] = "dense",
     ) -> GaussianDistribution:
         r"""Evaluate the Gaussian process posterior at the given points.
 
@@ -372,16 +434,28 @@ class AbstractPosterior(nnx.Module, tp.Generic[P, L]):
         Args:
             test_inputs: Input locations where the GP should be evaluated.
             train_data: Training dataset to condition on.
+            return_covariance_type: Literal denoting whether to return the full covariance
+                of the joint predictive distribution at the test_inputs (dense)
+                or just the the standard-deviation of the predictive distribution at
+                the test_inputs.
 
         Returns:
             GaussianDistribution: A multivariate normal random variable representation
                 of the Gaussian process.
         """
-        return self.predict(test_inputs, train_data)
+        return self.predict(
+            test_inputs,
+            train_data,
+            return_covariance_type=return_covariance_type,
+        )
 
     @abstractmethod
     def predict(
-        self, test_inputs: Num[Array, "N D"], train_data: Dataset
+        self,
+        test_inputs: Num[Array, "N D"],
+        train_data: Dataset,
+        *,
+        return_covariance_type: Literal["dense", "diagonal"] = "dense",
     ) -> GaussianDistribution:
         r"""Compute the latent function's multivariate normal distribution for a
         given set of parameters. For any class inheriting the `AbstractPosterior` class,
@@ -390,6 +464,10 @@ class AbstractPosterior(nnx.Module, tp.Generic[P, L]):
         Args:
             test_inputs: Input locations where the GP should be evaluated.
             train_data: Training dataset to condition on.
+            return_covariance_type: Literal denoting whether to return the full covariance
+                of the joint predictive distribution at the test_inputs (dense)
+                or just the the standard-deviation of the predictive distribution at
+                the test_inputs.
 
         Returns:
             GaussianDistribution: A multivariate normal random variable representation
@@ -442,8 +520,10 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
 
     def predict(
         self,
-        test_inputs: Num[Array, "N D"],
+        test_inputs: Num[Array, "M D"],
         train_data: Dataset,
+        *,
+        return_covariance_type: Literal["dense", "diagonal"] = "dense",
     ) -> GaussianDistribution:
         r"""Query the predictive posterior distribution.
 
@@ -454,13 +534,13 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
 
         The predictive distribution of a conjugate GP is given by
         $$
-            p(\mathbf{f}^{\star}\mid \mathbf{y}) & = \int p(\mathbf{f}^{\star} \mathbf{f} \mid \mathbf{y})\\
-            & =\mathcal{N}(\mathbf{f}^{\star} \boldsymbol{\mu}_{\mid \mathbf{y}}, \boldsymbol{\Sigma}_{\mid \mathbf{y}}
+        p(\mathbf{f}^{\star}\mid \mathbf{y}) & = \int p(\mathbf{f}^{\star} \mathbf{f} \mid \mathbf{y})\\
+        & =\mathcal{N}(\mathbf{f}^{\star} \boldsymbol{\mu}_{\mid \mathbf{y}}, \boldsymbol{\Sigma}_{\mid \mathbf{y}}
         $$
         where
         $$
-            \boldsymbol{\mu}_{\mid \mathbf{y}} & = k(\mathbf{x}^{\star}, \mathbf{x})\left(k(\mathbf{x}, \mathbf{x}')+\sigma^2\mathbf{I}_n\right)^{-1}\mathbf{y}  \\
-            \boldsymbol{\Sigma}_{\mid \mathbf{y}} & =k(\mathbf{x}^{\star}, \mathbf{x}^{\star\prime}) -k(\mathbf{x}^{\star}, \mathbf{x})\left( k(\mathbf{x}, \mathbf{x}') + \sigma^2\mathbf{I}_n \right)^{-1}k(\mathbf{x}, \mathbf{x}^{\star}).
+        \boldsymbol{\mu}_{\mid \mathbf{y}} & = k(\mathbf{x}^{\star}, \mathbf{x})\left(k(\mathbf{x}, \mathbf{x}')+\sigma^2\mathbf{I}_n\right)^{-1}\mathbf{y} \\
+        \boldsymbol{\Sigma}_{\mid \mathbf{y}} & =k(\mathbf{x}^{\star}, \mathbf{x}^{\star\prime}) -k(\mathbf{x}^{\star}, \mathbf{x})\left( k(\mathbf{x}, \mathbf{x}') + \sigma^2\mathbf{I}_n \right)^{-1}k(\mathbf{x}, \mathbf{x}^{\star}).
         $$
 
         The conditioning set is a GPJax `Dataset` object, whilst predictions
@@ -486,44 +566,65 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
                 predictive distribution is evaluated.
             train_data (Dataset): A `gpx.Dataset` object that contains the input and
                 output data used for training dataset.
+            return_covariance_type: Literal denoting whether to return the full covariance
+                of the joint predictive distribution at the test_inputs (dense)
+                or just the the standard-deviation of the predictive distribution at
+                the test_inputs.
 
         Returns:
             GaussianDistribution: A function that accepts an input array and
                 returns the predictive distribution as a `GaussianDistribution`.
         """
-        # Unpack training data
-        x, y = train_data.X, train_data.y
-
-        # Unpack test inputs
-        t = test_inputs
-
+        x = train_data.X
+        y = train_data.y
         # Observation noise o²
-        obs_noise = self.likelihood.obs_stddev.value**2
+        obs_noise = jnp.square(self.likelihood.obs_stddev.value)
         mx = self.prior.mean_function(x)
-
         # Precompute Gram matrix, Kxx, at training inputs, x
         Kxx = self.prior.kernel.gram(x)
-        Kxx_dense = add_jitter(Kxx.to_dense(), self.jitter)
-        Kxx = Dense(Kxx_dense)
+        Kxx = add_jitter(Kxx.to_dense(), self.jitter)
 
-        Sigma_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * obs_noise
+        Sigma_dense = Kxx + jnp.eye(Kxx.shape[0]) * obs_noise
         Sigma = psd(Dense(Sigma_dense))
         L_sigma = lower_cholesky(Sigma)
 
-        mean_t = self.prior.mean_function(t)
-        Ktt = self.prior.kernel.gram(t)
-        Kxt = self.prior.kernel.cross_covariance(x, t)
+        Kxt = self.prior.kernel.cross_covariance(x, test_inputs)
 
         L_inv_Kxt = solve(L_sigma, Kxt)
         L_inv_y_diff = solve(L_sigma, y - mx)
 
+        mean_t = self.prior.mean_function(test_inputs)
         mean = mean_t + jnp.matmul(L_inv_Kxt.T, L_inv_y_diff)
 
-        covariance = Ktt.to_dense() - jnp.matmul(L_inv_Kxt.T, L_inv_Kxt)
-        covariance = add_jitter(covariance, self.prior.jitter)
-        covariance = psd(Dense(covariance))
-
-        return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), covariance)
+        def _return_full_covariance(
+            L_inv_Kxt: Num[Array, "N M"],
+            t: Num[Array, "M D"],
+        ) -> Dense:
+            Ktt = self.prior.kernel.gram(t)
+            covariance = Ktt.to_dense() - jnp.matmul(L_inv_Kxt.T, L_inv_Kxt)
+            covariance = add_jitter(covariance, self.prior.jitter)
+            covariance = psd(Dense(covariance))
+            return covariance
+
+        def _return_diagonal_covariance(
+            L_inv_Kxt: Num[Array, "N M"],
+            t: Num[Array, "M D"],
+        ) -> Dense:
+            Ktt = self.prior.kernel.diagonal(t).diagonal
+            covariance = Ktt - jnp.einsum("ij, ji->i", L_inv_Kxt.T, L_inv_Kxt)
+            covariance += self.prior.jitter
+            covariance = psd(Dense(jnp.diag(jnp.atleast_1d(covariance.squeeze()))))
+            return covariance
+
+        cov = jax.lax.cond(
+            return_covariance_type == "dense",
+            _return_full_covariance,
+            _return_diagonal_covariance,
+            L_inv_Kxt,
+            test_inputs,
+        )
+
+        return GaussianDistribution(loc=jnp.atleast_1d(mean.squeeze()), scale=cov)
 
     def sample_approx(
         self,
@@ -567,7 +668,7 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
 
         Returns:
             FunctionalSample: A function representing an approximate sample from the Gaussian
-            process prior.
+                process prior.
         """
         if (not isinstance(num_samples, int)) or num_samples <= 0:
             raise ValueError("num_samples must be a positive integer")
@@ -586,7 +687,7 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
         canonical_weights = solve(
             Sigma,
             y + eps - jnp.inner(Phi, fourier_weights),
-        )  #  [N, B]
+        )  # [N, B]
 
         def sample_fn(test_inputs: Float[Array, "n D"]) -> Float[Array, "n B"]:
             fourier_features = fourier_feature_fn(test_inputs)
@@ -648,7 +749,11 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
         self.key = key
 
     def predict(
-        self, test_inputs: Num[Array, "N D"], train_data: Dataset
+        self,
+        test_inputs: Num[Array, "M D"],
+        train_data: Dataset,
+        *,
+        return_covariance_type: Literal["dense", "diagonal"] = "dense",
     ) -> GaussianDistribution:
         r"""Query the predictive posterior distribution.
 
@@ -660,50 +765,78 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
         transformed through the likelihood function's inverse link function.
 
         Args:
-            train_data (Dataset): A `gpx.Dataset` object that contains the input
-                and output data used for training dataset.
+        test_inputs (Num[Array, "N D"]): A Jax array of test inputs at which the
+            predictive distribution is evaluated.
+        train_data (Dataset): A `gpx.Dataset` object that contains the input
+            and output data used for training dataset.
+        return_covariance_type: Literal denoting whether to return the full covariance
+            of the joint predictive distribution at the test_inputs (dense)
+            or just the the standard-deviation of the predictive distribution at
+            the test_inputs.
 
         Returns:
             GaussianDistribution: A function that accepts an
                 input array and returns the predictive distribution as
                 a `dx.Distribution`.
         """
-        # Unpack training data
         x = train_data.X
-
-        # Unpack mean function and kernel
+        t = test_inputs
         mean_function = self.prior.mean_function
         kernel = self.prior.kernel
 
-        # Precompute lower triangular of Gram matrix, Lx, at training inputs, x
+        # Precompute lower triangular of Gram matrix
         Kxx = kernel.gram(x)
         Kxx_dense = add_jitter(Kxx.to_dense(), self.prior.jitter)
         Kxx = psd(Dense(Kxx_dense))
         Lx = lower_cholesky(Kxx)
 
-        # Unpack test inputs
-        t = test_inputs
-
-        # Compute terms of the posterior predictive distribution
-        Ktx = kernel.cross_covariance(t, x)
-        Ktt = kernel.gram(t)
-        mean_t = mean_function(t)
-
+        Kxt = kernel.cross_covariance(x, t)
         # Lx⁻¹ Kxt
-        Lx_inv_Kxt = solve(Lx, Ktx.T)
+        Lx_inv_Kxt = solve(Lx, Kxt)
 
+        mean_t = mean_function(t)
         # Whitened function values, wx, corresponding to the inputs, x
         wx = self.latent.value
 
         # μt + Ktx Lx⁻¹ wx
         mean = mean_t + jnp.matmul(Lx_inv_Kxt.T, wx)
 
-        # Ktt - Ktx Kxx⁻¹ Kxt, TODO: Take advantage of covariance structure to compute Schur complement more efficiently.
-        covariance = Ktt.to_dense() - jnp.matmul(Lx_inv_Kxt.T, Lx_inv_Kxt)
-        covariance = add_jitter(covariance, self.prior.jitter)
-        covariance = psd(Dense(covariance))
-
-        return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), covariance)
+        def _return_full_covariance(
+            Lx_inv_Kxt: Num[Array, "N M"],
+            t: Num[Array, "M D"],
+        ) -> Dense:
+            Ktt = kernel.gram(t)
+            covariance = Ktt.to_dense() - jnp.matmul(Lx_inv_Kxt.T, Lx_inv_Kxt)
+            covariance = add_jitter(covariance, self.prior.jitter)
+            covariance = psd(Dense(covariance))
+
+            return covariance
+
+        def _return_diagonal_covariance(
+            Lx_inv_Kxt: Num[Array, "N M"],
+            t: Num[Array, "M D"],
+        ) -> Dense:
+            Ktt = kernel.diagonal(t).diagonal
+            covariance = Ktt - jnp.einsum("ij, ji->i", Lx_inv_Kxt.T, Lx_inv_Kxt)
+            covariance += self.prior.jitter
+            # It would be nice to return a Diagonal here, but the pytree needs
+            # to be the same for both cond branches and the other branch needs
+            # to return a Dense.
+            # They are both LinearOperators, but they inherit from that class
+            # and hence are not the same pytree anymore.
+            covariance = psd(Dense(jnp.diag(jnp.atleast_1d(covariance.squeeze()))))
+
+            return covariance
+
+        cov = jax.lax.cond(
+            return_covariance_type == "dense",
+            _return_full_covariance,
+            _return_diagonal_covariance,
+            Lx_inv_Kxt,
+            test_inputs,
+        )
+
+        return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), cov)
 
 
 #######################

gpjax/kernels/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/base.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -123,7 +123,7 @@ class AbstractKernel(nnx.Module):
         """
         return self.compute_engine.gram(self, x)
 
-    def diagonal(self, x: Num[Array, "N D"]) -> Float[Array, " N"]:
+    def diagonal(self, x: Num[Array, "N D"]) -> LinearOperator:
         r"""Compute the diagonal of the gram matrix of the kernel.
 
         Args:

gpjax/kernels/computations/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/computations/base.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/computations/basis_functions.py

@@ -6,9 +6,7 @@ from jaxtyping import Float
 import gpjax
 from gpjax.kernels.computations.base import AbstractKernelComputation
 from gpjax.linalg import (
-    Dense,
     Diagonal,
-    psd,
 )
 from gpjax.typing import Array
 
@@ -27,9 +25,9 @@ class BasisFunctionComputation(AbstractKernelComputation):
         z2 = self.compute_features(kernel, y)
         return self.scaling(kernel) * jnp.matmul(z1, z2.T)
 
-    def _gram(self, kernel: K, inputs: Float[Array, "N D"]) -> Dense:
+    def _gram(self, kernel: K, inputs: Float[Array, "N D"]) -> Float[Array, "N N"]:
         z1 = self.compute_features(kernel, inputs)
-        return psd(Dense(self.scaling(kernel) * jnp.matmul(z1, z1.T)))
+        return self.scaling(kernel) * jnp.matmul(z1, z1.T)
 
     def diagonal(self, kernel: K, inputs: Float[Array, "N D"]) -> Diagonal:
         r"""For a given kernel, compute the elementwise diagonal of the

gpjax/kernels/computations/constant_diagonal.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/computations/dense.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/computations/diagonal.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/computations/eigen.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/non_euclidean/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/non_euclidean/graph.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,7 +17,7 @@ import beartype.typing as tp
 import jax.numpy as jnp
 from jaxtyping import (
     Float,
-    Int,
+    Integer,
     Num,
 )
 
@@ -103,11 +103,23 @@ class GraphKernel(StationaryKernel):
 
     def __call__(
         self,
-        x: Int[Array, "N 1"],
-        y: Int[Array, "M 1"],
+        x: ScalarInt | Integer[Array, " N"] | Integer[Array, "N 1"],
+        y: ScalarInt | Integer[Array, " M"] | Integer[Array, "M 1"],
     ):
+        x_idx = self._prepare_indices(x)
+        y_idx = self._prepare_indices(y)
         S = calculate_heat_semigroup(self)
-        Kxx = (jax_gather_nd(self.eigenvectors, x) * S.squeeze()) @ jnp.transpose(
-            jax_gather_nd(self.eigenvectors, y)
+        Kxx = (jax_gather_nd(self.eigenvectors, x_idx) * S.squeeze()) @ jnp.transpose(
+            jax_gather_nd(self.eigenvectors, y_idx)
         )  # shape (n,n)
         return Kxx.squeeze()
+
+    def _prepare_indices(
+        self,
+        indices: ScalarInt | Integer[Array, " N"] | Integer[Array, "N 1"],
+    ) -> Integer[Array, "N 1"]:
+        """Ensure index arrays are integer column vectors regardless of caller shape."""
+
+        idx = jnp.asarray(indices, dtype=jnp.int32)
+        idx = jnp.atleast_1d(idx)
+        return idx.reshape(-1, 1)

gpjax/kernels/non_euclidean/utils.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/nonstationary/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/nonstationary/arccosine.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/nonstationary/linear.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/nonstationary/polynomial.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/base.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/matern12.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/matern32.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/matern52.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/periodic.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/powered_exponential.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/rational_quadratic.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/rbf.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/utils.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/kernels/stationary/white.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2022 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

gpjax/scan.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The JaxGaussianProcesses Contributors. All Rights Reserved.
+# Copyright 2023 The thomaspinder Contributors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

{gpjax-0.13.1.dist-info → gpjax-0.13.3.dist-info}/METADATA

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.13.1
+Version: 0.13.3
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
-Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
-Project-URL: Source, https://github.com/JaxGaussianProcesses/GPJax
+Project-URL: Issues, https://github.com/thomaspinder/GPJax/issues
+Project-URL: Source, https://github.com/thomaspinder/GPJax
 Author-email: Thomas Pinder <tompinder@live.co.uk>
 License: MIT
 License-File: LICENSE.txt
@@ -72,11 +72,11 @@ Description-Content-Type: text/markdown
 <!-- <h1 align='center'>GPJax</h1>
 <h2 align='center'>Gaussian processes in Jax.</h2> -->
 <p align="center">
-<img width="700" height="300" src="https://raw.githubusercontent.com/JaxGaussianProcesses/GPJax/main/docs/static/gpjax_logo.svg" alt="GPJax's logo">
+<img width="700" height="300" src="https://raw.githubusercontent.com/thomaspinder/GPJax/main/docs/static/gpjax_logo.svg" alt="GPJax's logo">
 </p>
 
-[![codecov](https://codecov.io/gh/JaxGaussianProcesses/GPJax/branch/master/graph/badge.svg?token=DM1DRDASU2)](https://codecov.io/gh/JaxGaussianProcesses/GPJax)
-[![CodeFactor](https://www.codefactor.io/repository/github/jaxgaussianprocesses/gpjax/badge)](https://www.codefactor.io/repository/github/jaxgaussianprocesses/gpjax)
+[![codecov](https://codecov.io/gh/thomaspinder/GPJax/branch/master/graph/badge.svg?token=DM1DRDASU2)](https://codecov.io/gh/thomaspinder/GPJax)
+[![CodeFactor](https://www.codefactor.io/repository/github/thomaspinder/GPJax/badge)](https://www.codefactor.io/repository/github/thomaspinder/GPJax)
 [![Netlify Status](https://api.netlify.com/api/v1/badges/d3950e6f-321f-4508-9e52-426b5dae2715/deploy-status)](https://app.netlify.com/sites/endearing-crepe-c2d5fe/deploys)
 [![PyPI version](https://badge.fury.io/py/GPJax.svg)](https://badge.fury.io/py/GPJax)
 [![Conda Version](https://img.shields.io/conda/vn/conda-forge/gpjax.svg)](https://anaconda.org/conda-forge/gpjax)
@@ -101,19 +101,19 @@ with GP models.
 
 We would be delighted to receive contributions from interested individuals and
 groups. To learn how you can get involved, please read our [guide for
-contributing](https://github.com/JaxGaussianProcesses/GPJax/blob/main/docs/contributing.md).
+contributing](https://github.com/thomaspinder/GPJax/blob/main/docs/contributing.md).
 If you have any questions, we encourage you to [open an
-issue](https://github.com/JaxGaussianProcesses/GPJax/issues/new/choose). For
+issue](https://github.com/thomaspinder/GPJax/issues/new/choose). For
 broader conversations, such as best GP fitting practices or questions about the
 mathematics of GPs, we invite you to [open a
-discussion](https://github.com/JaxGaussianProcesses/GPJax/discussions).
+discussion](https://github.com/thomaspinder/GPJax/discussions).
 
 Another way you can contribute to GPJax is through [issue
 triaging](https://www.codetriage.com/what).  This can include reproducing bug reports,
 asking for vital information such as version numbers and reproduction instructions, or
 identifying stale issues. If you would like to begin triaging issues, an easy way to get
 started is to
-[subscribe to GPJax on CodeTriage](https://www.codetriage.com/jaxgaussianprocesses/gpjax).
+[subscribe to GPJax on CodeTriage](https://www.codetriage.com/thomaspinder/GPJax).
 
 As a contributor to GPJax, you are expected to abide by our [code of
 conduct](docs/CODE_OF_CONDUCT.md). If you feel that you have either experienced or
@@ -127,7 +127,7 @@ where we can discuss the development of GPJax and broader support for Gaussian
 process modelling.
 
 We appreciate all [the contributors to
-GPJax](https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors) who have helped to shape
+GPJax](https://github.com/thomaspinder/GPJax/graphs/contributors) who have helped to shape
 GPJax into the package it is today.
 
 # Supported methods and interfaces
@@ -215,7 +215,7 @@ conda install --channel conda-forge gpjax
 Clone a copy of the repository to your local machine and run the setup
 configuration in development mode.
 ```bash
-git clone https://github.com/JaxGaussianProcesses/GPJax.git
+git clone https://github.com/thomaspinder/GPJax.git
 cd GPJax
 uv venv
 uv sync --extra dev

gpjax-0.13.3.dist-info/RECORD

@@ -0,0 +1,52 @@
+gpjax/__init__.py,sha256=8EULPS_vtq4TeN6MjdLVHLlTVVyQADOoHtuVoRN8z5Y,1625
+gpjax/citation.py,sha256=pwFS8h1J-LE5ieRS0zDyuwhmQHNxkFHYE7iSMlVNmQc,3928
+gpjax/dataset.py,sha256=Ef5JGrl4jJS1mQmL3JdO0fdqbVmflT_Cu5VrlpYdJY4,4071
+gpjax/distributions.py,sha256=iKmeQ_NN2CIjRiuOeJlwEGASzGROi4ZCerVi1uY7zRM,7758
+gpjax/fit.py,sha256=tOXmM3l5-N3Jlnq8MOVkSpRj-0fRWOy2t9GxQRyUqxY,15318
+gpjax/gps.py,sha256=NcPXkkx0kXrSBTUje4QR6PS0DGrD7p5c-DyiOATwUz8,35338
+gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
+gpjax/likelihoods.py,sha256=xwnSQpn6Aa-FPpEoDn_3xpBdPQAmHP97jP-9iJmT4G8,9087
+gpjax/mean_functions.py,sha256=Aq09I5h5DZe1WokRxLa0Mpj8U0kJxpQ8CmdJJyCjNOc,6541
+gpjax/numpyro_extras.py,sha256=-vWJ7SpZVNhSdCjjrlxIkovMFrM1IzpsMJK3B4LioGE,3411
+gpjax/objectives.py,sha256=GvKbDIPqYjsc9FpiTccmZwRdHr2lCykgfxI9BX9I_GA,15362
+gpjax/parameters.py,sha256=hnyIKr6uIzd7Kb3KZC9WowR88ruQwUvdcto3cx2ZDv4,6756
+gpjax/scan.py,sha256=Z_V1yd76dL4B7_rnJnr_fohom6xzN_WRYxlTAAdqfa0,5357
+gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
+gpjax/variational_families.py,sha256=x4VnUh9GiW77ijnDEwrYzH0aWFOW6NLPQXEp--I9R-g,31566
+gpjax/kernels/__init__.py,sha256=GFiku1s8KPPzvfwDYzEjENPQReywvJ29M9Kl2JjkcoU,1893
+gpjax/kernels/base.py,sha256=gov60BgXOQK0PxbKQRVgjU55kCVZYAkDZrAInODvugo,11549
+gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
+gpjax/kernels/approximations/rff.py,sha256=WJ2Bw4RRW5bLXGKfStxph65H-ozDelvt_4mxXQm-g7w,4074
+gpjax/kernels/computations/__init__.py,sha256=gSkYASLcyPwcIg12Jjv7vxW9rixMVUhnlcl07iNvN_8,1371
+gpjax/kernels/computations/base.py,sha256=Fh_ppC7JwrbqKNJj1JqDOaUUGm23Zdm7RRLnGbQtePI,3608
+gpjax/kernels/computations/basis_functions.py,sha256=Y6OS4UEvo-Mdn7yRMWPde5rOvewDHIleaUIVBOpCTd4,2466
+gpjax/kernels/computations/constant_diagonal.py,sha256=ihxR0rcwcLi_sv3a026nEXRObMvmScBLOKj09XGwuns,1976
+gpjax/kernels/computations/dense.py,sha256=TGbAXSoaqo3Erzsaj7NztrEN2a_sHW5mi3bPbA1W9nc,1384
+gpjax/kernels/computations/diagonal.py,sha256=4V8aVXqC61GdgmFTvxDhmtMw1uSPi_C-m-d9-Fo6B8U,1760
+gpjax/kernels/computations/eigen.py,sha256=koVBKeLrHZX77PLlrykM-h2FN5g7l7S5xoyGUKD69Oc,1372
+gpjax/kernels/non_euclidean/__init__.py,sha256=Jb3K9a5-H-kNOirkrJJ46e3sgJIvoQS1-LfSsNXspJg,782
+gpjax/kernels/non_euclidean/graph.py,sha256=SevOCNeVZ5eK0MPhYGLnohAvNqig0C5c-oOpKuUKxeA,4631
+gpjax/kernels/non_euclidean/utils.py,sha256=PANh8fPOXnV8x6liTCJ08YYnDIQ33PrpLD-9hzjSg9w,2334
+gpjax/kernels/nonstationary/__init__.py,sha256=X6WVSMrejtLH6k_ct-1uZ9dmdmd_IEBIneweceXgXI0,922
+gpjax/kernels/nonstationary/arccosine.py,sha256=hjh6umWkSYs2xWa0rJQDPnmmkvk-vAwbGhMrkCev70o,5400
+gpjax/kernels/nonstationary/linear.py,sha256=xjeG4DN5K3-fgRw-HTFvPfcxYfYO5vPxdZc6K5qvnW8,2563
+gpjax/kernels/nonstationary/polynomial.py,sha256=D_czngVP0VGqFlqi_uJBayKegYs1Q-sHTCWCRDUWVAo,3380
+gpjax/kernels/stationary/__init__.py,sha256=b8KuCIENnW1nVQv3KdlZmOicfgIdoZzqgvVKlkgzkUk,1398
+gpjax/kernels/stationary/base.py,sha256=zIg21SK5nTg9kW0RJy6xXIp7XF7TjZxJT9Vvpsscz4M,6030
+gpjax/kernels/stationary/matern12.py,sha256=FbhdY5xlkgoCCOizVd7ZKK2XjJOCMDmfn2bTdxTOIKw,1763
+gpjax/kernels/stationary/matern32.py,sha256=tHHMN3pNMf5sGZNyCowi1RUUB0SENdkjHgQE8ZS3kMU,1956
+gpjax/kernels/stationary/matern52.py,sha256=TjS2_R9tYga8ePpSd5UIBTKDCsGWJX8AlUVvElYAKYk,2009
+gpjax/kernels/stationary/periodic.py,sha256=sCGWFYSTM6P91AFw9qTThAtzP5spatNrcU66Lr9qQjU,3520
+gpjax/kernels/stationary/powered_exponential.py,sha256=Z0q9B6mZysL8pKwFPf_kYsJ8hUOf2E8rSlxhbl0U3f8,3671
+gpjax/kernels/stationary/rational_quadratic.py,sha256=9X6X6YEIW8021XPzkCOMUihdWS0_JnBVO1b2JToFa2w,3410
+gpjax/kernels/stationary/rbf.py,sha256=uuhDlV-kZZ3gL0ZRpKWAOh3pVTpkddEVIILsWdEsMmM,1682
+gpjax/kernels/stationary/utils.py,sha256=wcS4rRJmNIeL3Log7gcK7klWWprulF4krQJq803alhk,2172
+gpjax/kernels/stationary/white.py,sha256=rPVf2xzfv2WpPpb8E7LXGuguiyZ6QKa7M3ltzLItmpk,2204
+gpjax/linalg/__init__.py,sha256=F8mxk_9Zc2nFd7Q-unjJ50_6rXEKzZj572WsU_jUKqI,547
+gpjax/linalg/operations.py,sha256=xvhOy5P4FmUCPWjIVNdg1yDXaoFQ48anFUfR-Tnfr6k,6480
+gpjax/linalg/operators.py,sha256=arxRGwcoAy_RqUYqBpZ3XG6OXbjShUl7m8sTpg85npE,11608
+gpjax/linalg/utils.py,sha256=fKV8G_iKZVhNkNvN20D_dQEi93-8xosGbXBP-v7UEyo,2020
+gpjax-0.13.3.dist-info/METADATA,sha256=1Au0xXIXWvS64SPLP8R_mPQ6nuUMayx2qdNLD-HdhW8,10296
+gpjax-0.13.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.13.3.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.13.3.dist-info/RECORD,,

gpjax-0.13.1.dist-info/RECORD

@@ -1,52 +0,0 @@
-gpjax/__init__.py,sha256=asMWra4r95NSlYQbniJhCQV6pEk39ONOTvnkm-wy8OA,1641
-gpjax/citation.py,sha256=pwFS8h1J-LE5ieRS0zDyuwhmQHNxkFHYE7iSMlVNmQc,3928
-gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
-gpjax/distributions.py,sha256=iKmeQ_NN2CIjRiuOeJlwEGASzGROi4ZCerVi1uY7zRM,7758
-gpjax/fit.py,sha256=I2sJVuKZii_d7MEcelHIivfM8ExYGMgdBuKKOT7Dw-A,15326
-gpjax/gps.py,sha256=fSbHjfQ1vKUZ3CnqdNgJVgFsuTkCwK3yK_c0SQsWYo0,30118
-gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
-gpjax/likelihoods.py,sha256=xwnSQpn6Aa-FPpEoDn_3xpBdPQAmHP97jP-9iJmT4G8,9087
-gpjax/mean_functions.py,sha256=Aq09I5h5DZe1WokRxLa0Mpj8U0kJxpQ8CmdJJyCjNOc,6541
-gpjax/numpyro_extras.py,sha256=-vWJ7SpZVNhSdCjjrlxIkovMFrM1IzpsMJK3B4LioGE,3411
-gpjax/objectives.py,sha256=GvKbDIPqYjsc9FpiTccmZwRdHr2lCykgfxI9BX9I_GA,15362
-gpjax/parameters.py,sha256=hnyIKr6uIzd7Kb3KZC9WowR88ruQwUvdcto3cx2ZDv4,6756
-gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
-gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
-gpjax/variational_families.py,sha256=x4VnUh9GiW77ijnDEwrYzH0aWFOW6NLPQXEp--I9R-g,31566
-gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
-gpjax/kernels/base.py,sha256=4vUV6_wnbV64EUaOecVWE9Vz4G5Dg6xUMNsQSis8A1o,11561
-gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
-gpjax/kernels/approximations/rff.py,sha256=WJ2Bw4RRW5bLXGKfStxph65H-ozDelvt_4mxXQm-g7w,4074
-gpjax/kernels/computations/__init__.py,sha256=uTVkqvnZVesFLDN92h0ZR0jfR69Eo2WyjOlmSYmCPJ8,1379
-gpjax/kernels/computations/base.py,sha256=L6K0roxZbrYeJKxEw-yaTiK9Mtcv0YtZfWI2Xnau7i8,3616
-gpjax/kernels/computations/basis_functions.py,sha256=_SFv4Tiwne40bxr1uVYpEjjZgjIQHKseLmss2Zgl1L4,2484
-gpjax/kernels/computations/constant_diagonal.py,sha256=JkQhLj7cK48IhOER4ivkALNhD1oQleKe-Rr9BtUJ6es,1984
-gpjax/kernels/computations/dense.py,sha256=vnW6XKQe4_gzpXRWTctxhgMA9-9TebdtiXzAqh_-j6g,1392
-gpjax/kernels/computations/diagonal.py,sha256=k1KqW0DwWRIBvbb7jzcKktXRfhXbcos3ncWrFplJ4W0,1768
-gpjax/kernels/computations/eigen.py,sha256=LuwYVPK0AuDGNSccPmAS8wyDJ2ngkRbc6kNZzqiJEOg,1380
-gpjax/kernels/non_euclidean/__init__.py,sha256=RT7puRPqCTpyxZ16q596EuOQEQi1LK1v3J9_fWz1NlY,790
-gpjax/kernels/non_euclidean/graph.py,sha256=MukL4ZghUywkau4tVWmHsrAjzTcl9kHN6onEAvxSvjc,4109
-gpjax/kernels/non_euclidean/utils.py,sha256=vxjRX209tOL3jKsw69hQfVtBqkASDZ__4tJOGI4oLjo,2342
-gpjax/kernels/nonstationary/__init__.py,sha256=YpWQfOy_cqOKc5ezn37vqoK3Z6jznYiJz28BD_8F7AY,930
-gpjax/kernels/nonstationary/arccosine.py,sha256=cqb8sqaNwW3fEbrA7MY9OF2KJFTkxHhqwmQtABE3G8w,5408
-gpjax/kernels/nonstationary/linear.py,sha256=UIMoCq2hg6dQKr4J5UGiiPqotBleQuYfy00Ia1NaMOo,2571
-gpjax/kernels/nonstationary/polynomial.py,sha256=CKc02C7Utgo-hhcOOCcKLdln5lj4vud_8M-JY7SevJ8,3388
-gpjax/kernels/stationary/__init__.py,sha256=j4BMTaQlIx2kNAT1Dkf4iO2rm-f7_oSVWNrk1bN0tqE,1406
-gpjax/kernels/stationary/base.py,sha256=YSgur73wqAFZBzt8D6CfujvyjowXoPOK1po10gjzMJo,6038
-gpjax/kernels/stationary/matern12.py,sha256=tfrIP-pJML-kSqVwf8wdVUwC53mhFohMnf7z5OMX3xE,1771
-gpjax/kernels/stationary/matern32.py,sha256=vdrszTtH03PB9D4AKg-YC_56xCupK3tO9L8vH0GYFcc,1964
-gpjax/kernels/stationary/matern52.py,sha256=IVsIKDo9bTZnA98yBmA_Q31peFdSCTXjf0_cVlLQh6k,2017
-gpjax/kernels/stationary/periodic.py,sha256=f4PhWhKg-pJsEBGzEMK9pdbylO84GPKhzHlBC83ZVWw,3528
-gpjax/kernels/stationary/powered_exponential.py,sha256=xuFGuIK0mKNMU3iLtZMXZTHXJuMFAMoX7gAtXefCdqU,3679
-gpjax/kernels/stationary/rational_quadratic.py,sha256=zHo2LVW65T52XET4Hx9JaKO0TfxylV8WRUtP7sUUOx0,3418
-gpjax/kernels/stationary/rbf.py,sha256=euHUs6FdfRICQcabAWE4MX-7GEDr2TxgZWdFQiXr9Bw,1690
-gpjax/kernels/stationary/utils.py,sha256=6BI9EBcCzeeKx-XH-MfW1ORmtU__tPX5zyvfLhpkBsU,2180
-gpjax/kernels/stationary/white.py,sha256=TkdXXZCCjDs7JwR_gj5uvn2s1wyfRbe1vyHhUMJ8jjI,2212
-gpjax/linalg/__init__.py,sha256=F8mxk_9Zc2nFd7Q-unjJ50_6rXEKzZj572WsU_jUKqI,547
-gpjax/linalg/operations.py,sha256=xvhOy5P4FmUCPWjIVNdg1yDXaoFQ48anFUfR-Tnfr6k,6480
-gpjax/linalg/operators.py,sha256=arxRGwcoAy_RqUYqBpZ3XG6OXbjShUl7m8sTpg85npE,11608
-gpjax/linalg/utils.py,sha256=fKV8G_iKZVhNkNvN20D_dQEi93-8xosGbXBP-v7UEyo,2020
-gpjax-0.13.1.dist-info/METADATA,sha256=qhsF8IgUu4oYTuLdqT56XQ4EKkDnUH_4D7imLxL_nPQ,10400
-gpjax-0.13.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-gpjax-0.13.1.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
-gpjax-0.13.1.dist-info/RECORD,,