gpjax 0.8.0__py3-none-any.whl → 0.8.2__py3-none-any.whl

gpjax/__init__.py

@@ -38,7 +38,7 @@ __license__ = "MIT"
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.8.0"
+__version__ = "0.8.2"
 
 __all__ = [
     "base",

gpjax/decision_making/test_functions/non_conjugate_functions.py

@@ -85,6 +85,6 @@ class PoissonTestFunction:
         Returns:
             Integer[Array, 'N 1']: Values of the test function at the points.
         """
-        key = jr.PRNGKey(42)
+        key = jr.key(42)
         f = lambda x: 2.0 * jnp.sin(3 * x) + 0.5 * x
         return jr.poisson(key, jnp.exp(f(x)))

gpjax/fit.py

@@ -23,11 +23,16 @@ from beartype.typing import (
     Union,
 )
 import jax
+from jax import (
+    jit,
+    value_and_grad,
+)
 from jax._src.random import _check_prng_key
+from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 import jax.random as jr
-import jaxopt
 import optax as ox
+import scipy
 
 from gpjax.base import Module
 from gpjax.dataset import Dataset
@@ -42,10 +47,6 @@ from gpjax.typing import (
 ModuleModel = TypeVar("ModuleModel", bound=Module)
 
 
-class FailedScipyFitError(Exception):
-    """Raised a model fit using Scipy fails"""
-
-
 def fit(  # noqa: PLR0913
     *,
     model: ModuleModel,
@@ -72,7 +73,7 @@ def fit(  # noqa: PLR0913
         >>>
         >>> # (1) Create a dataset:
         >>> X = jnp.linspace(0.0, 10.0, 100)[:, None]
-        >>> y = 2.0 * X + 1.0 + 10 * jr.normal(jr.PRNGKey(0), X.shape)
+        >>> y = 2.0 * X + 1.0 + 10 * jr.normal(jr.key(0), X.shape)
         >>> D = gpx.Dataset(X, y)
         >>>
         >>> # (2) Define your model:
@@ -110,7 +111,7 @@ def fit(  # noqa: PLR0913
         batch_size (Optional[int]): The size of the mini-batch to use. Defaults to -1
             (i.e. full batch).
         key (Optional[KeyArray]): The random key to use for the optimisation batch
-            selection. Defaults to jr.PRNGKey(42).
+            selection. Defaults to jr.key(42).
         log_rate (Optional[int]): How frequently the objective function's value should
             be printed. Defaults to 10.
         verbose (Optional[bool]): Whether to print the training loading bar. Defaults
@@ -130,7 +131,7 @@ def fit(  # noqa: PLR0913
         _check_optim(optim)
         _check_num_iters(num_iters)
         _check_batch_size(batch_size)
-        _check_prng_key(key)
+        _check_prng_key("fit", key)
         _check_log_rate(log_rate)
         _check_verbose(verbose)
 
@@ -214,30 +215,31 @@ def fit_scipy(  # noqa: PLR0913
     model = model.unconstrain()
 
     # Unconstrained space loss function with stop-gradient rule for non-trainable params.
-    def loss(model: Module, data: Dataset) -> ScalarFloat:
+    def loss(model: Module) -> ScalarFloat:
         model = model.stop_gradient()
-        return objective(model.constrain(), data)
-
-    solver = jaxopt.ScipyMinimize(
-        fun=loss,
-        maxiter=max_iters,
+        return objective(model.constrain(), train_data)
+
+    # convert to numpy for interface with scipy
+    x0, scipy_to_jnp = ravel_pytree(model)
+
+    @jit
+    def scipy_wrapper(x0):
+        value, grads = value_and_grad(loss)(scipy_to_jnp(jnp.array(x0)))
+        scipy_grads = ravel_pytree(grads)[0]
+        return value, scipy_grads
+
+    history = [scipy_wrapper(x0)[0]]
+    result = scipy.optimize.minimize(
+        fun=scipy_wrapper,
+        x0=x0,
+        jac=True,
+        callback=lambda X: history.append(scipy_wrapper(X)[0]),
+        options={"maxiter": max_iters, "disp": verbose},
     )
+    history = jnp.array(history)
 
-    initial_loss = solver.fun(model, train_data)
-    model, result = solver.run(model, data=train_data)
-    history = jnp.array([initial_loss, result.fun_val])
-
-    if verbose:
-        print(f"Initial loss is {initial_loss}")
-        if result.success:
-            print("Optimization was successful")
-        else:
-            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.
+    # convert back to pytree and reconstrain
+    model = scipy_to_jnp(result.x)
     model = model.constrain()
     return model, history
 

gpjax/flax_base/bijectors.py

@@ -0,0 +1,8 @@
+from gpjax.flax_base.types import BijectorLookupType, DomainType
+import tensorflow_probability.substrates.jax.bijectors as tfb
+import typing as tp
+
+Bijectors: BijectorLookupType = {
+    "real": tfb.Identity(),
+    "positive": tfb.Softplus(),
+}

gpjax/flax_base/param.py

@@ -0,0 +1,16 @@
+import jax.numpy as jnp
+import typing as tp
+from flax.experimental import nnx
+from gpjax.flax_base.types import DomainType, A
+
+
+class AbstractParameter(nnx.Variable[A]):
+    domain: DomainType = "real"
+    static: bool = False
+
+    def __init__(self, value: A, *args, **kwargs):
+        super().__init__(jnp.asarray(value), *args, **kwargs)
+
+
+class PositiveParameter(AbstractParameter[A]):
+    domain: DomainType = "positive"

gpjax/flax_base/types.py

@@ -0,0 +1,15 @@
+import typing as tp
+import tensorflow_probability.substrates.jax.bijectors as tfb
+
+DomainType = tp.Literal["real", "positive"]
+A = tp.TypeVar("A")
+
+
+# class BijectorLookup(tp.TypedDict):
+#     domain: DomainType
+#     bijector: tfb.Bijector
+
+class BijectorLookupType(tp.Dict[DomainType, tfb.Bijector]):
+    pass
+
+__all__ = ["DomainType", "A", "BijectorLookupType"]

gpjax/gps.py

@@ -12,10 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
+# from __future__ import annotations
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import overload
+from typing import (
+    TYPE_CHECKING,
+    Generic,
+    TypeVar,
+    overload,
+)
 
 from beartype.typing import (
     Any,
@@ -47,7 +52,7 @@ from gpjax.kernels.base import AbstractKernel
 from gpjax.likelihoods import (
     AbstractLikelihood,
     Gaussian,
-    NonGaussianLikelihood,
+    NonGaussian,
 )
 from gpjax.lower_cholesky import lower_cholesky
 from gpjax.mean_functions import AbstractMeanFunction
@@ -57,13 +62,19 @@ from gpjax.typing import (
     KeyArray,
 )
 
+Kernel = TypeVar("Kernel", bound=AbstractKernel)
+MeanFunction = TypeVar("MeanFunction", bound=AbstractMeanFunction)
+Likelihood = TypeVar("Likelihood", bound=AbstractLikelihood)
+NonGaussianLikelihood = TypeVar("NonGaussianLikelihood", bound=NonGaussian)
+GaussianLikelihood = TypeVar("GaussianLikelihood", bound=Gaussian)
+
 
 @dataclass
-class AbstractPrior(Module):
+class AbstractPrior(Module, Generic[MeanFunction, Kernel]):
     r"""Abstract Gaussian process prior."""
 
-    kernel: AbstractKernel
-    mean_function: AbstractMeanFunction
+    kernel: Kernel
+    mean_function: MeanFunction
     jitter: float = static_field(1e-6)
 
     def __call__(self, *args: Any, **kwargs: Any) -> GaussianDistribution:
@@ -113,7 +124,7 @@ class AbstractPrior(Module):
 # GP Priors
 #######################
 @dataclass
-class Prior(AbstractPrior):
+class Prior(AbstractPrior[MeanFunction, Kernel]):
     r"""A Gaussian process prior object.
 
     The GP is parameterised by a
@@ -137,17 +148,27 @@ class Prior(AbstractPrior):
     ```
     """
 
-    # @overload
-    # def __mul__(self, other: Gaussian) -> "ConjugatePosterior":
-    #     ...
+    if TYPE_CHECKING:
 
-    # @overload
-    # def __mul__(self, other: NonGaussianLikelihood) -> "NonConjugatePosterior":
-    #     ...
+        @overload
+        def __mul__(
+            self, other: GaussianLikelihood
+        ) -> "ConjugatePosterior[Prior[MeanFunction, Kernel], GaussianLikelihood]":
+            ...
 
-    # @overload
-    # def __mul__(self, other: AbstractLikelihood) -> "AbstractPosterior":
-    #     ...
+        @overload
+        def __mul__(
+            self, other: NonGaussianLikelihood
+        ) -> (
+            "NonConjugatePosterior[Prior[MeanFunction, Kernel], NonGaussianLikelihood]"
+        ):
+            ...
+
+        @overload
+        def __mul__(
+            self, other: Likelihood
+        ) -> "AbstractPosterior[Prior[MeanFunction, Kernel], Likelihood]":
+            ...
 
     def __mul__(self, other):
         r"""Combine the prior with a likelihood to form a posterior distribution.
@@ -183,17 +204,27 @@ class Prior(AbstractPrior):
         """
         return construct_posterior(prior=self, likelihood=other)
 
-    # @overload
-    # def __rmul__(self, other: Gaussian) -> "ConjugatePosterior":
-    #     ...
+    if TYPE_CHECKING:
+
+        @overload
+        def __rmul__(
+            self, other: GaussianLikelihood
+        ) -> "ConjugatePosterior[Prior[MeanFunction, Kernel], GaussianLikelihood]":
+            ...
 
-    # @overload
-    # def __rmul__(self, other: NonGaussianLikelihood) -> "NonConjugatePosterior":
-    #     ...
+        @overload
+        def __rmul__(
+            self, other: NonGaussianLikelihood
+        ) -> (
+            "NonConjugatePosterior[Prior[MeanFunction, Kernel], NonGaussianLikelihood]"
+        ):
+            ...
 
-    # @overload
-    # def __rmul__(self, other: AbstractLikelihood) -> "AbstractPosterior":
-    #     ...
+        @overload
+        def __rmul__(
+            self, other: Likelihood
+        ) -> "AbstractPosterior[Prior[MeanFunction, Kernel], Likelihood]":
+            ...
 
     def __rmul__(self, other):
         r"""Combine the prior with a likelihood to form a posterior distribution.
@@ -285,7 +316,7 @@ class Prior(AbstractPrior):
             >>> import gpjax as gpx
             >>> import jax.numpy as jnp
             >>> import jax.random as jr
-            >>> key = jr.PRNGKey(123)
+            >>> key = jr.key(123)
             >>>
             >>> meanf = gpx.mean_functions.Zero()
             >>> kernel = gpx.kernels.RBF()
@@ -324,19 +355,22 @@ class Prior(AbstractPrior):
         return sample_fn
 
 
+PriorType = TypeVar("PriorType", bound=AbstractPrior)
+
+
 #######################
 # GP Posteriors
 #######################
 @dataclass
-class AbstractPosterior(Module):
+class AbstractPosterior(Module, Generic[PriorType, Likelihood]):
     r"""Abstract Gaussian process posterior.
 
     The base GP posterior object conditioned on an observed dataset. All
     posterior objects should inherit from this class.
     """
 
-    prior: AbstractPrior
-    likelihood: AbstractLikelihood
+    prior: AbstractPrior[MeanFunction, Kernel]
+    likelihood: Likelihood
     jitter: float = static_field(1e-6)
 
     def __call__(self, *args: Any, **kwargs: Any) -> GaussianDistribution:
@@ -381,7 +415,7 @@ class AbstractPosterior(Module):
 
 
 @dataclass
-class ConjugatePosterior(AbstractPosterior):
+class ConjugatePosterior(AbstractPosterior[PriorType, GaussianLikelihood]):
     r"""A Conjuate Gaussian process posterior object.
 
     A Gaussian process posterior distribution when the constituent likelihood
@@ -600,7 +634,7 @@ class ConjugatePosterior(AbstractPosterior):
 
 
 @dataclass
-class NonConjugatePosterior(AbstractPosterior):
+class NonConjugatePosterior(AbstractPosterior[PriorType, NonGaussianLikelihood]):
     r"""A non-conjugate Gaussian process posterior object.
 
     A Gaussian process posterior object for models where the likelihood is
@@ -685,22 +719,17 @@ class NonConjugatePosterior(AbstractPosterior):
 #######################
 
 
-@overload
-def construct_posterior(prior: Prior, likelihood: Gaussian) -> ConjugatePosterior:
-    ...
-
-
 @overload
 def construct_posterior(
-    prior: Prior, likelihood: NonGaussianLikelihood
-) -> NonConjugatePosterior:
+    prior: PriorType, likelihood: GaussianLikelihood
+) -> ConjugatePosterior[PriorType, GaussianLikelihood]:
     ...
 
 
 @overload
 def construct_posterior(
-    prior: Prior, likelihood: AbstractLikelihood
-) -> AbstractPosterior:
+    prior: PriorType, likelihood: NonGaussianLikelihood
+) -> NonConjugatePosterior[PriorType, NonGaussianLikelihood]:
     ...
 
 

gpjax/likelihoods.py

@@ -247,11 +247,11 @@ def inv_probit(x: Float[Array, " *N"]) -> Float[Array, " *N"]:
     return 0.5 * (1.0 + jsp.special.erf(x / jnp.sqrt(2.0))) * (1 - 2 * jitter) + jitter
 
 
-NonGaussianLikelihood = Union[Poisson, Bernoulli]
+NonGaussian = Union[Poisson, Bernoulli]
 
 __all__ = [
     "AbstractLikelihood",
-    "NonGaussianLikelihood",
+    "NonGaussian",
     "Gaussian",
     "Bernoulli",
     "Poisson",

{gpjax-0.8.0.dist-info → gpjax-0.8.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpjax
-Version: 0.8.0
+Version: 0.8.2
 Summary: Gaussian processes in JAX.
 Home-page: https://github.com/JaxGaussianProcesses/GPJax
 License: Apache-2.0
@@ -14,14 +14,14 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Requires-Dist: beartype (>=0.16.2,<0.17.0)
 Requires-Dist: cola-ml (>=0.0.5,<0.0.6)
-Requires-Dist: jax (>=0.4.10)
-Requires-Dist: jaxlib (>=0.4.10)
-Requires-Dist: jaxopt (>=0.8.2,<0.9.0)
+Requires-Dist: jax (>=0.4.16)
+Requires-Dist: jaxlib (>=0.4.16)
+Requires-Dist: jaxopt (>=0.8.3,<0.9.0)
 Requires-Dist: jaxtyping (>=0.2.15,<0.3.0)
 Requires-Dist: optax (>=0.1.4,<0.2.0)
 Requires-Dist: orbax-checkpoint (>=0.2.3)
 Requires-Dist: simple-pytree (>=0.1.7,<0.2.0)
-Requires-Dist: tensorflow-probability (>=0.19.0,<0.20.0)
+Requires-Dist: tensorflow-probability (>=0.22.0,<0.23.0)
 Requires-Dist: tqdm (>=4.65.0,<5.0.0)
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Repository, https://github.com/JaxGaussianProcesses/GPJax
@@ -152,7 +152,7 @@ import jax.numpy as jnp
 import jax.random as jr
 import optax as ox
 
-key = jr.PRNGKey(123)
+key = jr.key(123)
 
 f = lambda x: 10 * jnp.sin(x)
 

{gpjax-0.8.0.dist-info → gpjax-0.8.2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
-gpjax/__init__.py,sha256=UPKdFOtBDHq93YqdBSar8kXQkvk2g-mE_IAgd-MbqMM,1545
+gpjax/__init__.py,sha256=HAAaTQX1KSXatt4Pm-aFLAYucVc3gg60OoMemZ81LrY,1545
 gpjax/base/__init__.py,sha256=jLAm6UYta8Uo02IQSsJGpzZI110XnQdyuAnHEX3A-2s,1068
 gpjax/base/module.py,sha256=50TgoIOKv58T13SwCoBpgU-3EH7e8piC9cuXMEIs5_c,12842
 gpjax/base/param.py,sha256=D_EQYfOf4wE1Xr3RBb5lyleeHYw8HKxNwU7Sr2LRmYw,2271
@@ -11,15 +11,18 @@ gpjax/decision_making/posterior_handler.py,sha256=fTI-mHl_WTrjXk0Pu_a91cFhDDZlWF
 gpjax/decision_making/search_space.py,sha256=3GRfM3KquOmOlLCzUvtELtB3HqJyaUOhxW9Jj2BnTtA,3547
 gpjax/decision_making/test_functions/__init__.py,sha256=GDCY9_kaAnxDWwzo1FkdxnDx-80MErAHchbGybT9xYs,1109
 gpjax/decision_making/test_functions/continuous_functions.py,sha256=ATrMPX_hE-uzgyjvsm1t_xzXmIkOrcGmTLgCQUOHUyw,5088
-gpjax/decision_making/test_functions/non_conjugate_functions.py,sha256=aDvnzZYmChY23nzOIETHrDoKvaJmZk81FMB0x9uYVfw,2986
+gpjax/decision_making/test_functions/non_conjugate_functions.py,sha256=MXFbPtRAZ7nFYs8eF3F9OBhm1Y4NrMs8w7b6uigAMGs,2982
 gpjax/decision_making/utility_functions/__init__.py,sha256=pAjgstn2l_HL0yMXP_vi6YfFX53eTdwAVIJeTA2gIlI,1171
 gpjax/decision_making/utility_functions/base.py,sha256=Vq50tPzmxzE_d9edRx_H85EN_XjaTFY3AkudUtXh9wY,4053
 gpjax/decision_making/utility_functions/thompson_sampling.py,sha256=vERF8aLKan_hxhLU6w5rGMLFcp1HGb3jrQzYuksg7bw,4344
 gpjax/decision_making/utility_maximizer.py,sha256=7J_qFwMwgfYxHU7ye9H7g78wHK9qdZvgW0orE0upJnM,6041
 gpjax/decision_making/utils.py,sha256=pIl4t5lhnWsfnAqrgUaQi6Ua82ogyKAezwhQIWo-4AE,1838
 gpjax/distributions.py,sha256=BwvpCSdcsi2kIFVLXKlMO7TWOUojVJ1K_7ai8EcCbq8,9520
-gpjax/fit.py,sha256=FGvtkYSYehOtscsZinAiqXCdNctFMtICD5ZqjnhDZeI,10654
-gpjax/gps.py,sha256=lPO_zuwf3rUpA2FAoVtAoGQ-NzItsUUbiW3MxPETRsA,29206
+gpjax/fit.py,sha256=kiGL9uSvYDcyVjmAPj8_QGorkRqzxMN_N03pO3U4RKY,10638
+gpjax/flax_base/bijectors.py,sha256=LpgZma8_fo9Kqi3JdRVse8CROoCG9PuNP-GTbaV63gY,244
+gpjax/flax_base/param.py,sha256=_ZkGO5o7vgBwm_l2suV2orfP7XIgk7Cj2D7_XXBxdDQ,426
+gpjax/flax_base/types.py,sha256=4_bGquCOuWNfOR-O6GLjZJJgueNwxu27Wckmm6Ckz_o,365
+gpjax/gps.py,sha256=lIgoGqxNI4AbZ80bjm_c_rNSNSy6KG3YkOlPJlE1gKQ,30328
 gpjax/integrators.py,sha256=wThbb-9d1FW1_33tg6Z7GLwruwuMm1-spd3-TEUCSAI,5763
 gpjax/kernels/__init__.py,sha256=pCTWcgf1hkJGYnX9BS-AcWryyZhovZ2Cfb4Ttf12J6E,1845
 gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
@@ -49,7 +52,7 @@ gpjax/kernels/stationary/rational_quadratic.py,sha256=JzrFdm5enVYdD61TtTnxyVrGBB
 gpjax/kernels/stationary/rbf.py,sha256=rIDwkF8YwOM2NfT7_y5Mfhc4BJSBOO6MdMeWPolspZ0,2442
 gpjax/kernels/stationary/utils.py,sha256=KU4j4IAXS1MZbHPUOiW6rieWX7lMCMRtSRlxIH3Kybo,2226
 gpjax/kernels/stationary/white.py,sha256=ymTKpOQZ-AWJqiZv4CbWt9pdpsw0HGfbf3jNIV83f_8,2086
-gpjax/likelihoods.py,sha256=j_l-AkgaYN_agnIehYYWc6ryUOxkoPp4CfftJ7nQmnc,8111
+gpjax/likelihoods.py,sha256=MPZntqagPsQ4Y00-KW9yyUvm9C2UHGX-pqN6GubxcCk,8091
 gpjax/lower_cholesky.py,sha256=4uM_2JF_fTDr4qq88N7mljR0MG5G7UVOmiqvvcf7hec,1850
 gpjax/mean_functions.py,sha256=6qxvaIoHfNMCYIN_9R_9HUiCTzBGiZzTk2vNlb4RaU8,6437
 gpjax/objectives.py,sha256=HKgbeGNUOm9kzndX7vXHYvHzTsE21GgAmD2VU9ktHzg,19178
@@ -57,7 +60,7 @@ gpjax/progress_bar.py,sha256=ErXoJ-jjEPOediq_clgG2OC0UTptpP0MtVM1a2ZLC_o,4418
 gpjax/scan.py,sha256=RTfIISMqo4jznDFJ6m9fiaPEYgnINRPwGNRiz71gLFg,5506
 gpjax/typing.py,sha256=H_7je0XuJ1UP77zP5bJskdK2gJhm812xRfnRbGdqgkg,1688
 gpjax/variational_families.py,sha256=5Q7btVizm3Y3xcVVpoosS-0JM2JreaQEiFbgF2uNgXc,26466
-gpjax-0.8.0.dist-info/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
-gpjax-0.8.0.dist-info/WHEEL,sha256=vxFmldFsRN_Hx10GDvsdv1wroKq8r5Lzvjp6GZ4OO8c,88
-gpjax-0.8.0.dist-info/METADATA,sha256=RJk6PlMtLj1QM5qP1lmpP38UEfbBOkRLJoo2NWgEaVg,10083
-gpjax-0.8.0.dist-info/RECORD,,
+gpjax-0.8.2.dist-info/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
+gpjax-0.8.2.dist-info/METADATA,sha256=yI9KlzNn3IATrcpijURB4MvKwXeXPfQVMpZcHmKpz5I,10079
+gpjax-0.8.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+gpjax-0.8.2.dist-info/RECORD,,

{gpjax-0.8.0.dist-info → gpjax-0.8.2.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.3.2
+Generator: poetry-core 1.9.0
 Root-Is-Purelib: true
 Tag: py3-none-any