def test_expand_trace_observations(self):
for batch_shape, dtype in itertools.product(
([], [2]), (torch.float, torch.double)
):
q, d = 3, 4
X = torch.rand(*batch_shape, q, d, device=self.device, dtype=dtype)
# test nullop behavior
self.assertTrue(torch.equal(expand_trace_observations(X), X))
self.assertTrue(
torch.equal(expand_trace_observations(X, fidelity_dims=[1]), X)
)
# test default behavior
num_tr = 2
X_expanded = expand_trace_observations(X, num_trace_obs=num_tr)
self.assertEqual(
X_expanded.shape, torch.Size(batch_shape + [q * (1 + num_tr), d])
)
for i in range(num_tr):
X_sub = X_expanded[..., q * i : q * (i + 1), :]
self.assertTrue(torch.equal(X_sub[..., :-1], X[..., :-1]))
X_sub_expected = (1 - i / (num_tr + 1)) * X[..., :q, -1]
self.assertTrue(torch.equal(X_sub[..., -1], X_sub_expected))
# test custom fidelity dims
fdims = [0, 2]
num_tr = 3
X_expanded = expand_trace_observations(
X, fidelity_dims=fdims, num_trace_obs=num_tr
)
self.assertEqual(
X_expanded.shape, torch.Size(batch_shape + [q * (1 + num_tr), d])
)
for j, i in itertools.product([1, 3], range(num_tr)):
X_sub = X_expanded[..., q * i : q * (i + 1), j]
self.assertTrue(torch.equal(X_sub, X[..., j]))
for j, i in itertools.product(fdims, range(num_tr)):
X_sub = X_expanded[..., q * i : q * (i + 1), j]
X_sub_expected = (1 - i / (1 + num_tr)) * X[..., :q, j]
self.assertTrue(torch.equal(X_sub, X_sub_expected))
# test gradients
num_tr = 2
fdims = [1]
X.requires_grad_(True)
X_expanded = expand_trace_observations(
X, fidelity_dims=fdims, num_trace_obs=num_tr
)
out = X_expanded.sum()
out.backward()
grad_exp = torch.full_like(X, 1 + num_tr)
grad_exp[..., fdims] = 1 + sum(
(i + 1) / (num_tr + 1) for i in range(num_tr)
)
self.assertTrue(torch.allclose(X.grad, grad_exp))
def expand(X: Tensor) -> Tensor:
return expand_trace_observations(
X=X,
fidelity_dims=sorted(target_fidelities),
# pyre-fixme[16]: `Optional` has no attribute `get`.
num_trace_obs=options.get(Keys.NUM_TRACE_OBSERVATIONS, 0),
)
def construct_inputs_mf_base(
model: Model,
training_data: TrainingData,
target_fidelities: Dict[int, Union[int, float]],
fidelity_weights: Optional[Dict[int, float]] = None,
cost_intercept: float = 1.0,
num_trace_observations: int = 0,
**ignore: Any,
) -> Dict[str, Any]:
r"""Construct kwargs for a multifidetlity acquisition function's constructor."""
if fidelity_weights is None:
fidelity_weights = {f: 1.0 for f in target_fidelities}
if 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_aware_utility = InverseCostWeightedUtility(
cost_model=AffineFidelityCostModel(
fidelity_weights=fidelity_weights, fixed_cost=cost_intercept
)
)
return {
"target_fidelities": target_fidelities,
"cost_aware_utility": cost_aware_utility,
"expand": lambda X: expand_trace_observations(
X=X,
fidelity_dims=sorted(target_fidelities),
num_trace_obs=num_trace_observations,
),
"project": lambda X: 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,
)
def test_construct_inputs_mf_base(self):
target_fidelities = {0: 0.123}
fidelity_weights = {0: 0.456}
cost_intercept = 0.789
num_trace_observations = 0
with self.subTest("test_fully_specified"):
kwargs = construct_inputs_mf_base(
model=mock.Mock(),
training_data=self.bd_td,
objective=LinearMCObjective(torch.rand(2)),
target_fidelities=target_fidelities,
fidelity_weights=fidelity_weights,
cost_intercept=cost_intercept,
num_trace_observations=num_trace_observations,
)
self.assertEqual(kwargs["target_fidelities"], target_fidelities)
X = torch.rand(3, 2)
self.assertTrue(isinstance(kwargs["expand"], Callable))
self.assertTrue(
torch.equal(
kwargs["expand"](X),
expand_trace_observations(
X=X,
fidelity_dims=sorted(target_fidelities),
num_trace_obs=num_trace_observations,
),
))
self.assertTrue(isinstance(kwargs["project"], Callable))
self.assertTrue(
torch.equal(
kwargs["project"](X),
project_to_target_fidelity(
X, target_fidelities=target_fidelities),
))
cm = kwargs["cost_aware_utility"].cost_model
w = torch.tensor(list(fidelity_weights.values()),
dtype=cm.weights.dtype)
self.assertEqual(cm.fixed_cost, cost_intercept)
self.assertTrue(torch.allclose(cm.weights, w))
with self.subTest("test_missing_fidelity_weights"):
kwargs = construct_inputs_mf_base(
model=mock.Mock(),
training_data=self.bd_td,
objective=LinearMCObjective(torch.rand(2)),
target_fidelities=target_fidelities,
cost_intercept=cost_intercept,
)
cm = kwargs["cost_aware_utility"].cost_model
self.assertTrue(
torch.allclose(cm.weights, torch.ones_like(cm.weights)))
with self.subTest("test_mismatched_weights"):
with self.assertRaisesRegex(RuntimeError,
"Must provide the same indices for"):
_ = construct_inputs_mf_base(
model=mock.Mock(),
training_data=self.bd_td,
objective=LinearMCObjective(torch.rand(2)),
target_fidelities={0: 1.0},
fidelity_weights={1: 0.5},
cost_intercept=cost_intercept,
)
