def test_evaluate_kg(self):
# a thorough test using real model and dtype double
d = 2
dtype = torch.double
bounds = torch.tensor([[0], [1]], device=self.device,
dtype=dtype).repeat(1, d)
train_X = torch.rand(3, d, device=self.device, dtype=dtype)
train_Y = torch.rand(3, 1, device=self.device, dtype=dtype)
model = SingleTaskGP(train_X, train_Y)
qKG = qKnowledgeGradient(
model=model,
num_fantasies=2,
objective=None,
X_pending=torch.rand(2, d, device=self.device, dtype=dtype),
current_value=torch.rand(1, device=self.device, dtype=dtype),
)
X = torch.rand(4, 3, d, device=self.device, dtype=dtype)
options = {"num_inner_restarts": 2, "raw_inner_samples": 3}
val = qKG.evaluate(X,
bounds=bounds,
num_restarts=2,
raw_samples=3,
options=options)
# verify output shape
self.assertEqual(val.size(), torch.Size([4]))
# verify dtype
self.assertEqual(val.dtype, dtype)
# test i) no dimension is squeezed out, ii) dtype float, iii) MC objective,
# and iv) t_batch_mode_transform
dtype = torch.float
bounds = torch.tensor([[0], [1]], device=self.device, dtype=dtype)
train_X = torch.rand(1, 1, device=self.device, dtype=dtype)
train_Y = torch.rand(1, 1, device=self.device, dtype=dtype)
model = SingleTaskGP(train_X, train_Y)
qKG = qKnowledgeGradient(
model=model,
num_fantasies=1,
objective=GenericMCObjective(
objective=lambda Y, X: Y.norm(dim=-1)),
)
X = torch.rand(1, 1, device=self.device, dtype=dtype)
options = {"num_inner_restarts": 1, "raw_inner_samples": 1}
val = qKG.evaluate(X,
bounds=bounds,
num_restarts=1,
raw_samples=1,
options=options)
# verify output shape
self.assertEqual(val.size(), torch.Size([1]))
# verify dtype
self.assertEqual(val.dtype, dtype)
def test_gen_one_shot_kg_initial_conditions(self):
num_fantasies = 8
num_restarts = 4
raw_samples = 16
for dtype in (torch.float, torch.double):
mean = torch.zeros(1, 1, device=self.device, dtype=dtype)
mm = MockModel(MockPosterior(mean=mean))
mock_kg = qKnowledgeGradient(model=mm, num_fantasies=num_fantasies)
bounds = torch.tensor([[0, 0], [1, 1]], device=self.device, dtype=dtype)
# test option error
with self.assertRaises(ValueError):
gen_one_shot_kg_initial_conditions(
acq_function=mock_kg,
bounds=bounds,
q=1,
num_restarts=num_restarts,
raw_samples=raw_samples,
options={"frac_random": 2.0},
)
# test generation logic
q = 2
mock_random_ics = torch.rand(num_restarts, q + num_fantasies, 2)
mock_fantasy_cands = torch.ones(20, 1, 2)
mock_fantasy_vals = torch.randn(20)
with ExitStack() as es:
mock_gbics = es.enter_context(
mock.patch(
"botorch.optim.initializers.gen_batch_initial_conditions",
return_value=mock_random_ics,
)
)
mock_optacqf = es.enter_context(
mock.patch(
"botorch.optim.optimize.optimize_acqf",
return_value=(mock_fantasy_cands, mock_fantasy_vals),
)
)
ics = gen_one_shot_kg_initial_conditions(
acq_function=mock_kg,
bounds=bounds,
q=q,
num_restarts=num_restarts,
raw_samples=raw_samples,
)
mock_gbics.assert_called_once()
mock_optacqf.assert_called_once()
n_value = int((1 - 0.1) * num_fantasies)
self.assertTrue(
torch.equal(
ics[..., :-n_value, :], mock_random_ics[..., :-n_value, :]
)
)
self.assertTrue(torch.all(ics[..., -n_value:, :] == 1))
def _instantiate_KG(
model: Model,
objective: AcquisitionObjective,
qmc: bool = True,
n_fantasies: int = 64,
mc_samples: int = 256,
num_trace_observations: int = 0,
seed_inner: Optional[int] = None,
seed_outer: Optional[int] = None,
X_pending: Optional[Tensor] = None,
current_value: Optional[Tensor] = None,
target_fidelities: Optional[Dict[int, float]] = None,
fidelity_weights: Optional[Dict[int, float]] = None,
cost_intercept: float = 1.0,
) -> qKnowledgeGradient:
r"""Instantiate either a `qKnowledgeGradient` or `qMultiFidelityKnowledgeGradient`
acquisition function depending on whether `target_fidelities` is defined.
"""
sampler_cls = SobolQMCNormalSampler if qmc else IIDNormalSampler
fantasy_sampler = sampler_cls(num_samples=n_fantasies, seed=seed_outer)
if isinstance(objective, MCAcquisitionObjective):
inner_sampler = sampler_cls(num_samples=mc_samples, seed=seed_inner)
else:
inner_sampler = None
if target_fidelities:
if fidelity_weights is None:
fidelity_weights = {f: 1.0 for f in target_fidelities}
if not set(target_fidelities) == set(fidelity_weights):
raise RuntimeError(
"Must provide the same indices for target_fidelities "
f"({set(target_fidelities)}) and fidelity_weights "
f" ({set(fidelity_weights)}).")
cost_model = AffineFidelityCostModel(fidelity_weights=fidelity_weights,
fixed_cost=cost_intercept)
cost_aware_utility = InverseCostWeightedUtility(cost_model=cost_model)
def project(X: Tensor) -> Tensor:
return project_to_target_fidelity(
X=X, target_fidelities=target_fidelities)
def expand(X: Tensor) -> Tensor:
return expand_trace_observations(
X=X,
fidelity_dims=sorted(target_fidelities), # pyre-ignore: [6]
num_trace_obs=num_trace_observations,
)
return qMultiFidelityKnowledgeGradient(
model=model,
num_fantasies=n_fantasies,
sampler=fantasy_sampler,
objective=objective,
inner_sampler=inner_sampler,
X_pending=X_pending,
current_value=current_value,
cost_aware_utility=cost_aware_utility,
project=project,
expand=expand,
)
return qKnowledgeGradient(
model=model,
num_fantasies=n_fantasies,
sampler=fantasy_sampler,
objective=objective,
inner_sampler=inner_sampler,
X_pending=X_pending,
current_value=current_value,
)
def test_initialize_q_knowledge_gradient(self):
for dtype in (torch.float, torch.double):
mean = torch.zeros(1, 1, device=self.device, dtype=dtype)
mm = MockModel(MockPosterior(mean=mean))
# test error when neither specifying neither sampler nor num_fantasies
with self.assertRaises(ValueError):
qKnowledgeGradient(model=mm, num_fantasies=None)
# test error when sampler and num_fantasies arg are inconsistent
sampler = IIDNormalSampler(num_samples=16)
with self.assertRaises(ValueError):
qKnowledgeGradient(model=mm, num_fantasies=32, sampler=sampler)
# test default construction
qKG = qKnowledgeGradient(model=mm, num_fantasies=32)
self.assertEqual(qKG.num_fantasies, 32)
self.assertIsInstance(qKG.sampler, SobolQMCNormalSampler)
self.assertEqual(qKG.sampler.sample_shape, torch.Size([32]))
self.assertIsNone(qKG.objective)
self.assertIsNone(qKG.inner_sampler)
self.assertIsNone(qKG.X_pending)
self.assertIsNone(qKG.current_value)
self.assertEqual(qKG.get_augmented_q_batch_size(q=3), 32 + 3)
# test custom construction
obj = GenericMCObjective(lambda Y, X: Y.mean(dim=-1))
sampler = IIDNormalSampler(num_samples=16)
X_pending = torch.zeros(2, 2, device=self.device, dtype=dtype)
qKG = qKnowledgeGradient(
model=mm,
num_fantasies=16,
sampler=sampler,
objective=obj,
X_pending=X_pending,
)
self.assertEqual(qKG.num_fantasies, 16)
self.assertEqual(qKG.sampler, sampler)
self.assertEqual(qKG.sampler.sample_shape, torch.Size([16]))
self.assertEqual(qKG.objective, obj)
self.assertIsInstance(qKG.inner_sampler, SobolQMCNormalSampler)
self.assertEqual(qKG.inner_sampler.sample_shape, torch.Size([128]))
self.assertTrue(torch.equal(qKG.X_pending, X_pending))
self.assertIsNone(qKG.current_value)
self.assertEqual(qKG.get_augmented_q_batch_size(q=3), 16 + 3)
# test assignment of num_fantasies from sampler if not provided
qKG = qKnowledgeGradient(model=mm, num_fantasies=None, sampler=sampler)
self.assertEqual(qKG.sampler.sample_shape, torch.Size([16]))
# test custom construction with inner sampler and current value
inner_sampler = SobolQMCNormalSampler(num_samples=256)
current_value = torch.zeros(1, device=self.device, dtype=dtype)
qKG = qKnowledgeGradient(
model=mm,
num_fantasies=8,
objective=obj,
inner_sampler=inner_sampler,
current_value=current_value,
)
self.assertEqual(qKG.num_fantasies, 8)
self.assertEqual(qKG.sampler.sample_shape, torch.Size([8]))
self.assertEqual(qKG.objective, obj)
self.assertIsInstance(qKG.inner_sampler, SobolQMCNormalSampler)
self.assertEqual(qKG.inner_sampler, inner_sampler)
self.assertIsNone(qKG.X_pending)
self.assertTrue(torch.equal(qKG.current_value, current_value))
self.assertEqual(qKG.get_augmented_q_batch_size(q=3), 8 + 3)
# test construction with non-MC objective (ScalarizedObjective)
qKG_s = qKnowledgeGradient(
model=mm,
num_fantasies=16,
sampler=sampler,
objective=ScalarizedObjective(weights=torch.rand(2)),
)
self.assertIsNone(qKG_s.inner_sampler)
self.assertIsInstance(qKG_s.objective, ScalarizedObjective)
# test error if no objective and multi-output model
mean2 = torch.zeros(1, 2, device=self.device, dtype=dtype)
mm2 = MockModel(MockPosterior(mean=mean2))
with self.assertRaises(UnsupportedError):
qKnowledgeGradient(model=mm2)
def test_evaluate_q_knowledge_gradient(self):
for dtype in (torch.float, torch.double):
# basic test
n_f = 4
mean = torch.rand(n_f, 1, 1, device=self.device, dtype=dtype)
variance = torch.rand(n_f, 1, 1, device=self.device, dtype=dtype)
mfm = MockModel(MockPosterior(mean=mean, variance=variance))
with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f:
with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs:
mock_num_outputs.return_value = 1
mm = MockModel(None)
qKG = qKnowledgeGradient(model=mm, num_fantasies=n_f)
X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype)
val = qKG(X)
patch_f.assert_called_once()
cargs, ckwargs = patch_f.call_args
self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1, 1]))
self.assertTrue(torch.allclose(val, mean.mean(), atol=1e-4))
self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
# batched evaluation
b = 2
mean = torch.rand(n_f, b, 1, device=self.device, dtype=dtype)
variance = torch.rand(n_f, b, 1, device=self.device, dtype=dtype)
mfm = MockModel(MockPosterior(mean=mean, variance=variance))
X = torch.rand(b, n_f + 1, 1, device=self.device, dtype=dtype)
with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f:
with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs:
mock_num_outputs.return_value = 1
mm = MockModel(None)
qKG = qKnowledgeGradient(model=mm, num_fantasies=n_f)
val = qKG(X)
patch_f.assert_called_once()
cargs, ckwargs = patch_f.call_args
self.assertEqual(ckwargs["X"].shape, torch.Size([b, 1, 1]))
self.assertTrue(
torch.allclose(val, mean.mean(dim=0).squeeze(-1), atol=1e-4)
)
self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
# pending points and current value
X_pending = torch.rand(2, 1, device=self.device, dtype=dtype)
mean = torch.rand(n_f, 1, 1, device=self.device, dtype=dtype)
variance = torch.rand(n_f, 1, 1, device=self.device, dtype=dtype)
mfm = MockModel(MockPosterior(mean=mean, variance=variance))
current_value = torch.rand(1, device=self.device, dtype=dtype)
X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype)
with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f:
with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs:
mock_num_outputs.return_value = 1
mm = MockModel(None)
qKG = qKnowledgeGradient(
model=mm,
num_fantasies=n_f,
X_pending=X_pending,
current_value=current_value,
)
val = qKG(X)
patch_f.assert_called_once()
cargs, ckwargs = patch_f.call_args
self.assertEqual(ckwargs["X"].shape, torch.Size([1, 3, 1]))
self.assertTrue(torch.allclose(val, mean.mean() - current_value, atol=1e-4))
self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
# test objective (inner MC sampling)
objective = GenericMCObjective(objective=lambda Y, X: Y.norm(dim=-1))
samples = torch.randn(3, 1, 1, device=self.device, dtype=dtype)
mfm = MockModel(MockPosterior(samples=samples))
X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype)
with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f:
with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs:
mock_num_outputs.return_value = 1
mm = MockModel(None)
qKG = qKnowledgeGradient(
model=mm, num_fantasies=n_f, objective=objective
)
val = qKG(X)
patch_f.assert_called_once()
cargs, ckwargs = patch_f.call_args
self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1, 1]))
self.assertTrue(torch.allclose(val, objective(samples).mean(), atol=1e-4))
self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
# test non-MC objective (ScalarizedObjective)
weights = torch.rand(2, device=self.device, dtype=dtype)
objective = ScalarizedObjective(weights=weights)
mean = torch.tensor([1.0, 0.5], device=self.device, dtype=dtype).expand(
n_f, 1, 2
)
cov = torch.tensor(
[[1.0, 0.1], [0.1, 0.5]], device=self.device, dtype=dtype
).expand(n_f, 2, 2)
posterior = GPyTorchPosterior(MultitaskMultivariateNormal(mean, cov))
mfm = MockModel(posterior)
with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f:
with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs:
mock_num_outputs.return_value = 2
mm = MockModel(None)
qKG = qKnowledgeGradient(
model=mm, num_fantasies=n_f, objective=objective
)
val = qKG(X)
patch_f.assert_called_once()
cargs, ckwargs = patch_f.call_args
self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1, 1]))
val_expected = (mean * weights).sum(-1).mean(0)
self.assertTrue(torch.allclose(val, val_expected))
def main(
benchmark_name,
dataset_name,
dimensions,
method_name,
num_runs,
run_start,
num_iterations,
acquisition_name,
# acquisition_optimizer_name,
gamma,
num_random_init,
mc_samples,
batch_size,
num_fantasies,
num_restarts,
raw_samples,
noise_variance_init,
# use_ard,
# use_input_warping,
standardize_targets,
input_dir,
output_dir):
# TODO(LT): Turn into options
# device = "cpu"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.double
benchmark = make_benchmark(benchmark_name,
dimensions=dimensions,
dataset_name=dataset_name,
input_dir=input_dir)
name = make_name(benchmark_name,
dimensions=dimensions,
dataset_name=dataset_name)
output_path = Path(output_dir).joinpath(name, method_name)
output_path.mkdir(parents=True, exist_ok=True)
options = dict(gamma=gamma,
num_random_init=num_random_init,
acquisition_name=acquisition_name,
mc_samples=mc_samples,
batch_size=batch_size,
num_restarts=num_restarts,
raw_samples=raw_samples,
num_fantasies=num_fantasies,
noise_variance_init=noise_variance_init,
standardize_targets=standardize_targets)
with output_path.joinpath("options.yaml").open('w') as f:
yaml.dump(options, f)
config_space = DenseConfigurationSpace(benchmark.get_config_space())
bounds = create_bounds(config_space.get_bounds(),
device=device,
dtype=dtype)
input_dim = config_space.get_dimensions()
def func(tensor, *args, **kwargs):
"""
Wrapper that receives and returns torch.Tensor
"""
config = dict_from_tensor(tensor, cs=config_space)
# turn into maximization problem
res = -benchmark.evaluate(config).value
return torch.tensor(res, device=device, dtype=dtype)
for run_id in trange(run_start, num_runs, unit="run"):
run_begin_t = batch_end_t_adj = batch_end_t = datetime.now()
frames = []
features = []
targets = []
noise_variance = torch.tensor(noise_variance_init,
device=device,
dtype=dtype)
state_dict = None
with trange(num_iterations) as iterations:
for batch in iterations:
if len(targets) < num_random_init:
# click.echo(f"Completed {i}/{num_random_init} initial runs. "
# "Suggesting random candidate...")
# TODO(LT): support random seed
X_batch = torch.rand(size=(batch_size, input_dim),
device=device,
dtype=dtype)
else:
# construct dataset
X = torch.vstack(features)
y = torch.hstack(targets).unsqueeze(axis=-1)
y = standardize(y) if standardize_targets else y
# construct model
# model = FixedNoiseGP(X, standardize(y), noise_variance.expand_as(y),
model = FixedNoiseGP(X,
y,
noise_variance.expand_as(y),
input_transform=None).to(X)
mll = ExactMarginalLogLikelihood(model.likelihood, model)
if state_dict is not None:
model.load_state_dict(state_dict)
# update model
fit_gpytorch_model(mll)
# construct acquisition function
tau = torch.quantile(y, q=1 - gamma)
iterations.set_postfix(tau=tau.item())
if acquisition_name == "q-KG":
assert num_fantasies is not None and num_fantasies > 0
acq = qKnowledgeGradient(model,
num_fantasies=num_fantasies)
elif acquisition_name == "q-EI":
assert mc_samples is not None and mc_samples > 0
qmc_sampler = SobolQMCNormalSampler(
num_samples=mc_samples)
acq = qExpectedImprovement(model=model,
best_f=tau,
sampler=qmc_sampler)
# optimize acquisition function
X_batch, b = optimize_acqf(acq_function=acq,
bounds=bounds,
q=batch_size,
num_restarts=num_restarts,
raw_samples=raw_samples,
options=dict(batch_limit=5,
maxiter=200))
state_dict = model.state_dict()
# begin batch evaluation
batch_begin_t = datetime.now()
decision_duration = batch_begin_t - batch_end_t
batch_begin_t_adj = batch_end_t_adj + decision_duration
eval_end_times = []
# TODO(LT): Deliberately not doing broadcasting for now since
# batch sizes are so small anyway. Can revisit later if there
# is a compelling reason to do it.
rows = []
for j, x_next in enumerate(X_batch):
# eval begin time
eval_begin_t = datetime.now()
# evaluate blackbox objective
y_next = func(x_next)
# eval end time
eval_end_t = datetime.now()
# eval duration
eval_duration = eval_end_t - eval_begin_t
# adjusted eval end time is the duration added to the
# time at which batch eval was started
eval_end_t_adj = batch_begin_t_adj + eval_duration
eval_end_times.append(eval_end_t_adj)
elapsed = eval_end_t_adj - run_begin_t
# update dataset
features.append(x_next)
targets.append(y_next)
row = dict_from_tensor(x_next, cs=config_space)
row["loss"] = -y_next.item()
row["cost_eval"] = eval_duration.total_seconds()
row["finished"] = elapsed.total_seconds()
rows.append(row)
batch_end_t = datetime.now()
batch_end_t_adj = max(eval_end_times)
frame = pd.DataFrame(data=rows) \
.assign(batch=batch,
cost_decision=decision_duration.total_seconds())
frames.append(frame)
data = pd.concat(frames, axis="index", ignore_index=True)
data.to_csv(output_path.joinpath(f"{run_id:03d}.csv"))
return 0
def test_evaluate_q_knowledge_gradient(self):
for dtype in (torch.float, torch.double):
# basic test
n_f = 4
mean = torch.rand(n_f, 1, device=self.device, dtype=dtype)
variance = torch.rand(n_f, 1, device=self.device, dtype=dtype)
mfm = MockModel(MockPosterior(mean=mean, variance=variance))
with mock.patch.object(MockModel, "fantasize",
return_value=mfm) as patch_f:
mm = MockModel(None)
qKG = qKnowledgeGradient(model=mm, num_fantasies=n_f)
X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype)
val = qKG(X)
patch_f.assert_called_once()
cargs, ckwargs = patch_f.call_args
self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1]))
self.assertTrue(torch.allclose(val, mean.mean(), atol=1e-4))
self.assertTrue(
torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
# batched evaluation
b = 2
mean = torch.rand(n_f, b, 1, device=self.device, dtype=dtype)
variance = torch.rand(n_f, b, 1, device=self.device, dtype=dtype)
mfm = MockModel(MockPosterior(mean=mean, variance=variance))
X = torch.rand(b, n_f + 1, 1, device=self.device, dtype=dtype)
with mock.patch.object(MockModel, "fantasize",
return_value=mfm) as patch_f:
mm = MockModel(None)
qKG = qKnowledgeGradient(model=mm, num_fantasies=n_f)
val = qKG(X)
patch_f.assert_called_once()
cargs, ckwargs = patch_f.call_args
self.assertEqual(ckwargs["X"].shape, torch.Size([b, 1, 1]))
self.assertTrue(
torch.allclose(val, mean.mean(dim=0).squeeze(-1), atol=1e-4))
self.assertTrue(
torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
# pending points and current value
mean = torch.rand(n_f, 1, device=self.device, dtype=dtype)
variance = torch.rand(n_f, 1, device=self.device, dtype=dtype)
X_pending = torch.rand(2, 1, device=self.device, dtype=dtype)
mfm = MockModel(MockPosterior(mean=mean, variance=variance))
current_value = torch.rand(1, device=self.device, dtype=dtype)
X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype)
with mock.patch.object(MockModel, "fantasize",
return_value=mfm) as patch_f:
mm = MockModel(None)
qKG = qKnowledgeGradient(
model=mm,
num_fantasies=n_f,
X_pending=X_pending,
current_value=current_value,
)
val = qKG(X)
patch_f.assert_called_once()
cargs, ckwargs = patch_f.call_args
self.assertEqual(ckwargs["X"].shape, torch.Size([3, 1]))
self.assertTrue(
torch.allclose(val, mean.mean() - current_value, atol=1e-4))
self.assertTrue(
torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
# test objective (inner MC sampling)
objective = GenericMCObjective(objective=lambda Y: Y.norm(dim=-1))
samples = torch.randn(3, 1, 1, device=self.device, dtype=dtype)
mfm = MockModel(MockPosterior(samples=samples))
X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype)
with mock.patch.object(MockModel, "fantasize",
return_value=mfm) as patch_f:
mm = MockModel(None)
qKG = qKnowledgeGradient(model=mm,
num_fantasies=n_f,
objective=objective)
val = qKG(X)
patch_f.assert_called_once()
cargs, ckwargs = patch_f.call_args
self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1]))
self.assertTrue(
torch.allclose(val, objective(samples).mean(), atol=1e-4))
self.assertTrue(
torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
