gpjax 0.9.4__py3-none-any.whl → 0.10.0__py3-none-any.whl

gpjax/__init__.py

@@ -19,7 +19,6 @@ from beartype.roar import BeartypeDecorHintPep585DeprecationWarning
 filterwarnings("ignore", category=BeartypeDecorHintPep585DeprecationWarning)
 
 from gpjax import (
-    decision_making,
     gps,
     integrators,
     kernels,
@@ -40,11 +39,10 @@ __license__ = "MIT"
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.9.4"
+__version__ = "0.10.0"
 
 __all__ = [
     "base",
-    "decision_making",
     "gps",
     "integrators",
     "kernels",

gpjax/citation.py

@@ -10,11 +10,6 @@ from beartype.typing import (
 )
 from jaxlib.xla_extension import PjitFunction
 
-from gpjax.decision_making.test_functions import (
-    Forrester,
-    LogarithmicGoldsteinPrice,
-)
-from gpjax.decision_making.utility_functions import ThompsonSampling
 from gpjax.kernels import (
     RFF,
     ArcCosine,
@@ -149,41 +144,3 @@ def _(tree) -> PaperCitation:
         booktitle="Advances in neural information processing systems",
         citation_type="article",
     )
-
-
-####################
-# Decision making citations
-####################
-@cite.register(ThompsonSampling)
-def _(tree) -> PaperCitation:
-    return PaperCitation(
-        citation_key="wilson2020efficiently",
-        title="Efficiently sampling functions from Gaussian process posteriors",
-        authors="Wilson, James and Borovitskiy, Viacheslav and Terenin, Alexander and Mostowsky, Peter and Deisenroth, Marc",
-        year="2020",
-        booktitle="International Conference on Machine Learning",
-        citation_type="article",
-    )
-
-
-@cite.register(Forrester)
-def _(tree) -> BookCitation:
-    return BookCitation(
-        citation_key="forrester2008engineering",
-        authors="Forrester, Alexander and Sobester, Andras and Keane, Andy",
-        title="Engineering design via surrogate modelling: a practical guide",
-        year="2008",
-        publisher="John Wiley & Sons",
-    )
-
-
-@cite.register(LogarithmicGoldsteinPrice)
-def _(tree) -> PaperCitation:
-    return PaperCitation(
-        citation_key="picheny2013benchmark",
-        authors="Picheny, Victor and Wagner, Tobias and Ginsbourger, David",
-        title="A benchmark of kriging-based infill criteria for noisy optimization",
-        year="2013",
-        booktitle="Structural and multidisciplinary optimization",
-        citation_type="article",
-    )

gpjax/distributions.py

@@ -162,7 +162,9 @@ class GaussianDistribution(tfd.Distribution):
 
         return vmap(affine_transformation)(Z)
 
-    def sample(self, seed: KeyArray, sample_shape: Tuple[int, ...]):  # pylint: disable=useless-super-delegation
+    def sample(
+        self, seed: KeyArray, sample_shape: Tuple[int, ...]
+    ):  # pylint: disable=useless-super-delegation
         r"""See `Distribution.sample`."""
         return self._sample_n(
             seed, sample_shape[0]

{gpjax-0.9.4.dist-info → gpjax-0.10.0.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.9.4
+Version: 0.10.0
 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
 Author-email: Thomas Pinder <tompinder@live.co.uk>
-License-Expression: Apache-2.0
-License-File: LICENSE
+License: MIT
+License-File: LICENSE.txt
 Keywords: gaussian-processes jax machine-learning bayesian
 Classifier: Development Status :: 4 - Beta
 Classifier: Programming Language :: Python
@@ -18,13 +18,12 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: <3.13,>=3.10
 Requires-Dist: beartype>0.16.1
-Requires-Dist: cola-ml==0.0.5
-Requires-Dist: flax<0.10.0
-Requires-Dist: jax<0.4.28
-Requires-Dist: jaxlib<0.4.28
-Requires-Dist: jaxopt==0.8.2
+Requires-Dist: cola-ml>=0.0.7
+Requires-Dist: flax>=0.10.0
+Requires-Dist: jax>=0.5.0
+Requires-Dist: jaxlib>=0.5.0
 Requires-Dist: jaxtyping>0.2.10
-Requires-Dist: numpy<2.0.0
+Requires-Dist: numpy>=2.0.0
 Requires-Dist: optax>0.2.1
 Requires-Dist: tensorflow-probability>=0.24.0
 Requires-Dist: tqdm>4.66.2

{gpjax-0.9.4.dist-info → gpjax-0.10.0.dist-info}/RECORD

@@ -1,7 +1,7 @@
-gpjax/__init__.py,sha256=f1Sl-8Oz6YuEueKxvzIAL0iH_9b9xGzQv07tddS5wto,1697
-gpjax/citation.py,sha256=R4Pmvjt0ndA0avEDSvIbxDxKapkRRYXWX7RRWBvZCRQ,5306
+gpjax/__init__.py,sha256=LeAdMRx9XYvLf6csLhCIv6IHnDbAFB9rP--TYXECgz0,1654
+gpjax/citation.py,sha256=f2Hzj5MLyCE7l0hHAzsEQoTORZH5hgV_eis4uoBiWvE,3811
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
-gpjax/distributions.py,sha256=zxkSEZIlTg0PHvvgj0BQuIFEg-ugx6_NkEwSsbqWUM0,9325
+gpjax/distributions.py,sha256=X48FJr3reop9maherdMVt7-XZOm2f26T8AJt_IKM_oE,9339
 gpjax/fit.py,sha256=OHv8jUHxa1ndpqMERSDRtYtUDzubk9rMPVIhfCiIH5Q,11551
 gpjax/gps.py,sha256=97lYGrsmsufQxKEd8qz5wPNvui6FKXTF_Ps-sMFIjnY,31246
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
@@ -13,20 +13,6 @@ gpjax/parameters.py,sha256=Z4Wy3gEzPZG23-dtqC437_ZWnd_sPe9LcLCKn21ZBvA,4886
 gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
 gpjax/variational_families.py,sha256=s1rk7PtNTjQPabmVu-jBsuJBoqsxAAXwKFZJOEswkNQ,28161
-gpjax/decision_making/__init__.py,sha256=SDuPQl80lJ7nhfRsiB_7c22wCMiQO5ehSNohxUGnB7w,2170
-gpjax/decision_making/decision_maker.py,sha256=S4pOXrWcEHy0NDA0gfWzhk7pG0NJfaPpMXvq03yTy0g,13915
-gpjax/decision_making/posterior_handler.py,sha256=UgXf1Gu7GMh2YDSmiSWJIzmWlFW06KTS44HYz3mazZQ,5905
-gpjax/decision_making/search_space.py,sha256=bXwtMOhHZ2klnABpXm5Raxe7b0NTRDjo_cN3ecbk53Y,3545
-gpjax/decision_making/utility_maximizer.py,sha256=VT2amwSJbB64IL_MiWNl9ZgjcqO757qK6NW2gUBKsqs,5965
-gpjax/decision_making/utils.py,sha256=5j1GO5kcmG2laZR39NjhqgEjRekAWWzrnREv_5Zct_Y,2367
-gpjax/decision_making/test_functions/__init__.py,sha256=GDCY9_kaAnxDWwzo1FkdxnDx-80MErAHchbGybT9xYs,1109
-gpjax/decision_making/test_functions/continuous_functions.py,sha256=oL5ZQkvmbC3u9rEvSYI2DRAN3r7Ynf7wRZQlUWjKjt0,5612
-gpjax/decision_making/test_functions/non_conjugate_functions.py,sha256=eJpCnTS9dRieLxpjH4L6OTsP-w9JM3XhjnzCfk2Xqn8,2957
-gpjax/decision_making/utility_functions/__init__.py,sha256=xXI-4JKWAfTJ7XZ1vRDpqtb91MNzSPD0lP6xo0tOc7o,1445
-gpjax/decision_making/utility_functions/base.py,sha256=FOqrsRDmtHiCVl6IHr12-AEYBLStzMT5EBs-F92e1Og,3882
-gpjax/decision_making/utility_functions/expected_improvement.py,sha256=H6hjC-lj1oiHf2BomeQqroORQ7vtcOngiDAWxRwkNbg,4481
-gpjax/decision_making/utility_functions/probability_of_improvement.py,sha256=O_rHH1yR34JJlpAueSDJ_yo95fPI2aAGkwphS8snBYk,5220
-gpjax/decision_making/utility_functions/thompson_sampling.py,sha256=S-Yyn-9jsKkaXTvKFBP4sG_eCCKApGbHao5RR5tqXAo,4353
 gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
 gpjax/kernels/base.py,sha256=abkj3zidsBs7YSkYEfjeJ5jTs1YyDCPoBM2ZzqaqrgI,11561
 gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
@@ -56,7 +42,7 @@ gpjax/kernels/stationary/rational_quadratic.py,sha256=dYONp3i4rnKj3ET8UyxAKXv6UO
 gpjax/kernels/stationary/rbf.py,sha256=G13gg5phO7ite7D9QgoCy7gB2_y0FM6GZhgFW4RL6Xw,1734
 gpjax/kernels/stationary/utils.py,sha256=Xa9EEnxgFqEi08ZSFAZYYHhJ85_3Ac-ZUyUk18B63M4,2225
 gpjax/kernels/stationary/white.py,sha256=TkdXXZCCjDs7JwR_gj5uvn2s1wyfRbe1vyHhUMJ8jjI,2212
-gpjax-0.9.4.dist-info/METADATA,sha256=Qx_Qv91sE7_Y-c9CGuF40QBFk2FjLW0Fo2SHqFAgQFQ,10010
-gpjax-0.9.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-gpjax-0.9.4.dist-info/licenses/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
-gpjax-0.9.4.dist-info/RECORD,,
+gpjax-0.10.0.dist-info/METADATA,sha256=wZyZSD1p2t_K5m25TrvGJr6lTlfaUxoB12F-0f1d9Co,9970
+gpjax-0.10.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.10.0.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.10.0.dist-info/RECORD,,

gpjax-0.10.0.dist-info/licenses/LICENSE.txt

@@ -0,0 +1,19 @@
+(C) Copyright 2019 Hewlett Packard Enterprise Development LP
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation
+the rights to use, copy, modify, merge, publish, distribute, sublicense,
+and/or sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included
+in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
+OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+OTHER DEALINGS IN THE SOFTWARE.

gpjax/decision_making/__init__.py

@@ -1,63 +0,0 @@
-# Copyright 2023 The JaxGaussianProcesses 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.
-# You may obtain a copy of the License at
-#
-#     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,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-from gpjax.decision_making.decision_maker import (
-    AbstractDecisionMaker,
-    UtilityDrivenDecisionMaker,
-)
-from gpjax.decision_making.posterior_handler import PosteriorHandler
-from gpjax.decision_making.search_space import (
-    AbstractSearchSpace,
-    ContinuousSearchSpace,
-)
-from gpjax.decision_making.test_functions import (
-    AbstractContinuousTestFunction,
-    Forrester,
-    LogarithmicGoldsteinPrice,
-    Quadratic,
-)
-from gpjax.decision_making.utility_functions import (
-    AbstractSinglePointUtilityFunctionBuilder,
-    AbstractUtilityFunctionBuilder,
-    SinglePointUtilityFunction,
-    ThompsonSampling,
-    UtilityFunction,
-)
-from gpjax.decision_making.utility_maximizer import (
-    AbstractSinglePointUtilityMaximizer,
-    AbstractUtilityMaximizer,
-    ContinuousSinglePointUtilityMaximizer,
-)
-from gpjax.decision_making.utils import build_function_evaluator
-
-__all__ = [
-    "AbstractUtilityFunctionBuilder",
-    "AbstractUtilityMaximizer",
-    "AbstractDecisionMaker",
-    "AbstractSearchSpace",
-    "AbstractSinglePointUtilityFunctionBuilder",
-    "AbstractSinglePointUtilityMaximizer",
-    "UtilityFunction",
-    "build_function_evaluator",
-    "ContinuousSinglePointUtilityMaximizer",
-    "ContinuousSearchSpace",
-    "UtilityDrivenDecisionMaker",
-    "AbstractContinuousTestFunction",
-    "Forrester",
-    "LogarithmicGoldsteinPrice",
-    "PosteriorHandler",
-    "Quadratic",
-    "SinglePointUtilityFunction",
-    "ThompsonSampling",
-]

gpjax/decision_making/decision_maker.py

@@ -1,302 +0,0 @@
-# Copyright 2023 The GPJax 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.
-# You may obtain a copy of the License at
-#
-#     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,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-from abc import (
-    ABC,
-    abstractmethod,
-)
-import copy
-from dataclasses import dataclass
-
-from beartype.typing import (
-    Callable,
-    Dict,
-    List,
-    Mapping,
-)
-import jax.numpy as jnp
-import jax.random as jr
-
-from gpjax.dataset import Dataset
-from gpjax.decision_making.posterior_handler import PosteriorHandler
-from gpjax.decision_making.search_space import AbstractSearchSpace
-from gpjax.decision_making.utility_functions import (
-    AbstractUtilityFunctionBuilder,
-    ThompsonSampling,
-)
-from gpjax.decision_making.utility_maximizer import AbstractUtilityMaximizer
-from gpjax.decision_making.utils import FunctionEvaluator
-from gpjax.gps import AbstractPosterior
-from gpjax.typing import (
-    Array,
-    Float,
-    KeyArray,
-)
-
-
-@dataclass
-class AbstractDecisionMaker(ABC):
-    """
-    AbstractDecisionMaker abstract base class which handles the core decision making
-    loop, where we sequentially decide on points to query our function of interest at.
-    The decision making loop is split into two key steps, `ask` and `tell`. The `ask`
-    step is typically used to decide which point to query next. The `tell` step is
-    typically used to update models and datasets with newly queried points. These steps
-    can be combined in a 'run' loop which alternates between asking which point to query
-    next and telling the decision maker about the newly queried point having evaluated
-    the black-box function of interest at this point.
-
-    Attributes:
-        search_space: Search space over which we can evaluate the function(s) of interest.
-        posterior_handlers: dictionary of posterior handlers, which are used to update
-            posteriors throughout the decision making loop. Note that the word `posteriors`
-            is used for consistency with GPJax, but these objects are typically referred to
-            as `models` in the model-based decision making literature. Tags are used to
-            distinguish between posteriors. In a typical Bayesian optimisation setup one of
-            the tags will be `OBJECTIVE`, defined in `decision_making.utils`.
-        datasets: dictionary of datasets, which are augmented with
-            observations throughout the decision making loop. In a typical setup they are
-            also used to update the posteriors, using the `posterior_handlers`. Tags are used
-            to distinguish datasets, and correspond to tags in `posterior_handlers`.
-        key: JAX random key, used to generate random numbers.
-        batch_size: Number of points to query at each step of the decision making
-            loop. Note that `SinglePointUtilityFunction`s are only capable of generating
-            one point to be queried at each iteration of the decision making loop.
-        post_ask: List of functions to be executed after each ask step.
-        post_tell: List of functions to be executed after each tell step.
-    """
-
-    search_space: AbstractSearchSpace
-    posterior_handlers: Dict[str, PosteriorHandler]
-    datasets: Dict[str, Dataset]
-    key: KeyArray
-    batch_size: int
-    post_ask: List[
-        Callable
-    ]  # Specific type is List[Callable[[AbstractDecisionMaker, Float[Array, ["B D"]]], None]] but causes Beartype issues
-    post_tell: List[
-        Callable
-    ]  # Specific type is List[Callable[[AbstractDecisionMaker], None]] but causes Beartype issues
-
-    def __post_init__(self):
-        """
-        At initialisation we check that the posterior handlers and datasets are
-        consistent (i.e. have the same tags), and then initialise the posteriors, optimizing them using the
-        corresponding datasets.
-        """
-        self.datasets = copy.copy(
-            self.datasets
-        )  # Ensure initial datasets passed in to DecisionMaker are not mutated from within
-
-        if self.batch_size < 1:
-            raise ValueError(
-                f"Batch size must be greater than 0, got {self.batch_size}."
-            )
-
-        # Check that posterior handlers and datasets are consistent
-        if self.posterior_handlers.keys() != self.datasets.keys():
-            raise ValueError(
-                "Posterior handlers and datasets must have the same keys. "
-                f"Got posterior handlers keys {self.posterior_handlers.keys()} and "
-                f"datasets keys {self.datasets.keys()}."
-            )
-
-        # Initialize posteriors
-        self.posteriors: Dict[str, AbstractPosterior] = {}
-        for tag, posterior_handler in self.posterior_handlers.items():
-            self.posteriors[tag] = posterior_handler.get_posterior(
-                self.datasets[tag], optimize=True, key=self.key
-            )
-
-    @abstractmethod
-    def ask(self, key: KeyArray) -> Float[Array, "B D"]:
-        """
-        Get the point(s) to be queried next.
-
-        Args:
-            key (KeyArray): JAX PRNG key for controlling random state.
-
-        Returns:
-            Float[Array, "1 D"]: Point to be queried next
-        """
-        raise NotImplementedError
-
-    def tell(self, observation_datasets: Mapping[str, Dataset], key: KeyArray):
-        """
-        Add newly observed data to datasets and update the corresponding posteriors.
-
-        Args:
-            observation_datasets: dictionary of datasets containing new observations.
-                Tags are used to distinguish datasets, and correspond to tags in
-                `posterior_handlers` and `self.datasets`.
-            key: JAX PRNG key for controlling random state.
-        """
-        if observation_datasets.keys() != self.datasets.keys():
-            raise ValueError(
-                "Observation datasets and existing datasets must have the same keys. "
-                f"Got observation datasets keys {observation_datasets.keys()} and "
-                f"existing datasets keys {self.datasets.keys()}."
-            )
-
-        for tag, observation_dataset in observation_datasets.items():
-            self.datasets[tag] += observation_dataset
-
-        for tag, posterior_handler in self.posterior_handlers.items():
-            key, _ = jr.split(key)
-            self.posteriors[tag] = posterior_handler.update_posterior(
-                self.datasets[tag], self.posteriors[tag], optimize=True, key=key
-            )
-
-    def run(
-        self, n_steps: int, black_box_function_evaluator: FunctionEvaluator
-    ) -> Mapping[str, Dataset]:
-        """
-        Run the decision making loop continuously for for `n_steps`. This is broken down
-        into three main steps:
-        1. Call the `ask` method to get the point to be queried next.
-        2. Call the `black_box_function_evaluator` to evaluate the black box functions
-        of interest at the point chosen to be queried.
-        3. Call the `tell` method to update the datasets and posteriors with the newly
-        observed data.
-
-        In addition to this, after the `ask` step, the functions in the `post_ask` list
-        are executed, taking as arguments the decision maker and the point chosen to be
-        queried next. Similarly, after the `tell` step, the functions in the `post_tell`
-        list are executed, taking the decision maker as the sole argument.
-
-        Args:
-            n_steps (int): Number of steps to run the decision making loop for.
-            black_box_function_evaluator (FunctionEvaluator): Function evaluator which
-                evaluates the black box functions of interest at supplied points.
-
-        Returns:
-            Mapping[str, Dataset]: Dictionary of datasets containing the observations
-            made throughout the decision making loop, as well as the initial data
-            supplied when initialising the `DecisionMaker`.
-        """
-        for _ in range(n_steps):
-            query_point = self.ask(self.key)
-
-            for post_ask_method in self.post_ask:
-                post_ask_method(self, query_point)
-
-            self.key, _ = jr.split(self.key)
-            observation_datasets = black_box_function_evaluator(query_point)
-            self.tell(observation_datasets, self.key)
-
-            for post_tell_method in self.post_tell:
-                post_tell_method(self)
-
-        return self.datasets
-
-
-@dataclass
-class UtilityDrivenDecisionMaker(AbstractDecisionMaker):
-    """
-    UtilityDrivenDecisionMaker class which handles the core decision making loop in a
-    typical model-based decision making setup. In this setup we use surrogate model(s)
-    for the function(s) of interest, and define a utility function (often called the
-    'acquisition function' in the context of Bayesian optimisation) which characterises
-    how useful it would be to query a given point within the search space given the data
-    we have observed so far. This can then be used to decide which point(s) to query
-    next.
-
-    The decision making loop is split into two key steps, `ask` and `tell`. The `ask`
-    step forms a `UtilityFunction` from the current `posteriors` and `datasets` and
-    returns the point which maximises it. It also stores the formed utility function
-    under the attribute `self.current_utility_function` so that it can be called,
-    for instance for plotting, after the `ask` function has been called. The `tell` step
-    adds a newly queried point to the `datasets` and updates the `posteriors`.
-
-    This can be run as a typical ask-tell loop, or the `run` method can be used to run
-    the decision making loop for a fixed number of steps. Moreover, the `run` method executes
-    the functions in `post_ask` and `post_tell` after each ask and tell step
-    respectively. This enables the user to add custom functionality, such as the ability
-    to plot values of interest during the optimization process.
-
-    Attributes:
-        utility_function_builder (AbstractUtilityFunctionBuilder): Object which
-                builds utility functions from posteriors and datasets, to decide where
-                to query next. In a typical Bayesian optimisation setup the point chosen to
-                be queried next is the point which maximizes the utility function.
-        utility_maximizer (AbstractUtilityMaximizer): Object which maximizes
-            utility functions over the search space.
-    """
-
-    utility_function_builder: AbstractUtilityFunctionBuilder
-    utility_maximizer: AbstractUtilityMaximizer
-
-    def __post_init__(self):
-        super().__post_init__()
-        if self.batch_size > 1 and not isinstance(
-            self.utility_function_builder, ThompsonSampling
-        ):
-            raise NotImplementedError(
-                "Batch size > 1 currently only supported for Thompson sampling."
-            )
-
-    def ask(self, key: KeyArray) -> Float[Array, "B D"]:
-        """
-        Get updated utility function(s) and return the point(s) which maximises it/them. This
-        method also stores the utility function(s) in
-        `self.current_utility_functions` so that they can be accessed after the ask
-        function has been called. This is useful for non-deterministic utility
-        functions, which may differ between calls to `ask` due to the splitting of
-        `self.key`.
-
-        Note that in general `SinglePointUtilityFunction`s are only capable of
-        generating one point to be queried at each iteration of the decision making loop
-        (i.e. `self.batch_size` must be 1). However, Thompson sampling can be used in a
-        batched setting by drawing a batch of different samples from the GP posterior.
-        This is done by calling `build_utility_function` with different keys
-        sequentilly, and optimising each of these individual samples in sequence in
-        order to obtain `self.batch_size` points to query next.
-
-        Args:
-            key (KeyArray): JAX PRNG key for controlling random state.
-
-        Returns:
-            Float[Array, "B D"]: Point(s) to be queried next.
-        """
-        self.current_utility_functions = []
-        maximizers = []
-        # We currently only allow Thompson sampling to be run with batch size > 1. More
-        # batched utility functions may be added in the future.
-        if isinstance(self.utility_function_builder, ThompsonSampling) or (
-            (not isinstance(self.utility_function_builder, ThompsonSampling))
-            and (self.batch_size == 1)
-        ):
-            # Draw 'self.batch_size' Thompson samples and optimize each of them in order to
-            # obtain 'self.batch_size' points to query next.
-            for _ in range(self.batch_size):
-                decision_function = (
-                    self.utility_function_builder.build_utility_function(
-                        self.posteriors, self.datasets, key
-                    )
-                )
-                self.current_utility_functions.append(decision_function)
-
-                _, key = jr.split(key)
-                maximizer = self.utility_maximizer.maximize(
-                    decision_function, self.search_space, key
-                )
-                maximizers.append(maximizer)
-                _, key = jr.split(key)
-
-            maximizers = jnp.concatenate(maximizers)
-            return maximizers
-        else:
-            raise NotImplementedError(
-                "Only Thompson sampling currently supports batch size > 1."
-            )