def prepare_acquisition_function(
self,
dataset: Dataset,
model: ProbabilisticModel,
pending_points: TensorType = None,
) -> AcquisitionFunction:
if pending_points is None:
return lambda at: -tf.reduce_max(model.predict(at)[0], axis=-2)
else:
best_pending_score = tf.reduce_max(
model.predict(pending_points)[0])
return lambda at: -tf.math.maximum(
model.predict(at)[0][0], best_pending_score)
def _prepare_acquisition_function_assert(
_: object, dataset: Dataset, model: ProbabilisticModel
) -> Callable[[TensorType], TensorType]:
npt.assert_allclose(dataset.query_points, 0.0)
_, var = model.predict(tf.constant([0.0]))
npt.assert_allclose(var, 0.0)
return raise_
def _build_with_sampler(
self,
dataset: Dataset,
model: ProbabilisticModel,
sampler: IndependentReparametrizationSampler,
) -> AcquisitionFunction:
x = tf.random.uniform([100, 2], minval=-10.0, maxval=10.0, dtype=tf.float64)
samples = sampler.sample(x)
mean, var = model.predict(x)
_assert_kolmogorov_smirnov_95(
tf.linalg.matrix_transpose(samples),
tfp.distributions.Normal(mean[..., None], tf.sqrt(var)[..., None]),
)
return raise_