Prepare for future batching #385
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
In the future certain utility functions will be "batched" utility functions which quantify the utility of querying a *batch* of points rather than a single point. These utility functions may in turn be optimised differently. Therefore, I propose we have a `SinglePointUtilityFunction` type for non-batched utility functions, which has type annotation `[N D] -> [N 1]`, and in the future we may add `BatchUtilityFunction` which has type annotation `[N B D] -> [N 1]`. Distinguishing between these makes it easy for users to see which utility functions support batching. Also refactored UtilityDrivenDecisionMaker to enable for Thompson sampling to be used in a batched context, and updated unit tests to reflect this. Finally, added a new notebook to serve as an introduction to the `decision_making` module.
thomaspinder
approved these changes
Sep 7, 2023
henrymoss
reviewed
Sep 8, 2023
Comment on lines
+261
to
+268
| 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. | ||
|
|
henrymoss
reviewed
Sep 8, 2023
| maximizers = jnp.concatenate(maximizers) | ||
| return maximizers | ||
| else: | ||
| raise NotImplementedError( |
There was a problem hiding this comment.
this keeps it tidy for now
henrymoss
reviewed
Sep 8, 2023
| maximizers = [] | ||
| if isinstance(self.utility_function_builder, ThompsonSampling) or ( | ||
| (not isinstance(self.utility_function_builder, ThompsonSampling)) | ||
| and (self.batch_size == 1) |
There was a problem hiding this comment.
does this still not work with batch_size=1? This logic is a little tricky to understand
There was a problem hiding this comment.
It's restricting batch size to be 1 for non-TS acquisition functions if someone happens to implement one for now.
There was a problem hiding this comment.
got it. Maybe you can make the logic clearer so people like me dont get confused?
Thomas-Christie
merged commit b225ecf
into
JaxGaussianProcesses:tchristie/bo
20 checks passed
8 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Type of changes
Checklist
poetry run pre-commit run --all-files --show-diff-on-failurebefore committing.Description
In the future certain utility functions will be "batched" utility functions which quantify the utility of querying a batch of points rather than a single point. These utility functions may in turn be optimised differently.
Therefore, I propose we have a
SinglePointUtilityFunctiontype for non-batched utility functions, which has type annotation[N D] -> [N 1], and in the future we may addBatchUtilityFunctionwhich has type annotation[N B D] -> [N 1]. Distinguishing between these makes it easy for users to see which utility functions support batching.Also refactored UtilityDrivenDecisionMaker to enable for Thompson sampling to be used in a batched context, and updated unit tests to reflect this. Finally, added a new notebook to serve as an introduction to the
decision_makingmodule.Issue Number: N/A