def draw_proposals(
bounds,
originals: ep.Tensor,
perturbed: ep.Tensor,
unnormalized_source_directions: ep.Tensor,
source_directions: ep.Tensor,
source_norms: ep.Tensor,
spherical_steps: ep.Tensor,
source_steps: ep.Tensor,
):
# remember the actual shape
shape = originals.shape
assert perturbed.shape == shape
assert unnormalized_source_directions.shape == shape
assert source_directions.shape == shape
# flatten everything to (batch, size)
originals = flatten(originals)
perturbed = flatten(perturbed)
unnormalized_source_directions = flatten(unnormalized_source_directions)
source_directions = flatten(source_directions)
N, D = originals.shape
assert source_norms.shape == (N, )
assert spherical_steps.shape == (N, )
assert source_steps.shape == (N, )
# draw from an iid Gaussian (we can share this across the whole batch)
eta = ep.normal(perturbed, (D, 1))
# make orthogonal (source_directions are normalized)
eta = eta.T - ep.matmul(source_directions, eta) * source_directions
assert eta.shape == (N, D)
# rescale
norms = l2norms(eta)
assert norms.shape == (N, )
eta = eta * atleast_kd(spherical_steps * source_norms / norms, eta.ndim)
# project on the sphere using Pythagoras
distances = atleast_kd((spherical_steps.square() + 1).sqrt(), eta.ndim)
directions = eta - unnormalized_source_directions
spherical_candidates = originals + directions / distances
# clip
min_, max_ = bounds
spherical_candidates = spherical_candidates.clip(min_, max_)
# step towards the original inputs
new_source_directions = originals - spherical_candidates
assert new_source_directions.ndim == 2
new_source_directions_norms = l2norms(new_source_directions)
# length if spherical_candidates would be exactly on the sphere
lengths = source_steps * source_norms
# length including correction for numerical deviation from sphere
lengths = lengths + new_source_directions_norms - source_norms
# make sure the step size is positive
lengths = ep.maximum(lengths, 0)
# normalize the length
lengths = lengths / new_source_directions_norms
lengths = atleast_kd(lengths, new_source_directions.ndim)
candidates = spherical_candidates + lengths * new_source_directions
# clip
candidates = candidates.clip(min_, max_)
# restore shape
candidates = candidates.reshape(shape)
spherical_candidates = spherical_candidates.reshape(shape)
return candidates, spherical_candidates
def f(x: Tensor) -> Tuple[Tensor, Tuple[Tensor, Tensor]]:
x = x.square()
return x.sum(), (x, x + 1)
def f(x: ep.Tensor) -> ep.Tensor:
return x.square().sum()
def f(x: Tensor) -> Tensor:
return x.square().sum()
def draw_proposals(bounds: Bounds, originals: ep.Tensor, perturbed: ep.Tensor,
unnormalized_source_directions: ep.Tensor,
source_directions: ep.Tensor, source_norms: ep.Tensor,
spherical_steps: ep.Tensor, source_steps: ep.Tensor,
surrogate_model: Model) -> Tuple[ep.Tensor, ep.Tensor]:
# remember the actual shape
shape = originals.shape
assert perturbed.shape == shape
assert unnormalized_source_directions.shape == shape
assert source_directions.shape == shape
# flatten everything to (batch, size)
originals = flatten(originals)
perturbed = flatten(perturbed)
unnormalized_source_directions = flatten(unnormalized_source_directions)
source_directions = flatten(source_directions)
N, D = originals.shape
assert source_norms.shape == (N, )
assert spherical_steps.shape == (N, )
assert source_steps.shape == (N, )
# draw from an iid Gaussian (we can share this across the whole batch)
eta = ep.normal(perturbed, (D, 1))
# make orthogonal (source_directions are normalized)
eta = eta.T - ep.matmul(source_directions, eta) * source_directions
assert eta.shape == (N, D)
pg_factor = 0.5
if not surrogate_model is None:
device = surrogate_model.device
projected_gradient = get_projected_gradients(perturbed.reshape(shape),
originals.reshape(shape),
0, surrogate_model)
projected_gradient = projected_gradient.reshape((N, D))
projected_gradient = torch.tensor(projected_gradient, device=device)
projected_gradient, restore_type = ep.astensor_(projected_gradient)
eta = (1. - pg_factor) * eta + pg_factor * projected_gradient
# rescale
norms = ep.norms.l2(eta, axis=-1)
assert norms.shape == (N, )
eta = eta * atleast_kd(spherical_steps * source_norms / norms, eta.ndim)
# project on the sphere using Pythagoras
distances = atleast_kd((spherical_steps.square() + 1).sqrt(), eta.ndim)
directions = eta - unnormalized_source_directions
spherical_candidates = originals + directions / distances
# clip
min_, max_ = bounds
spherical_candidates = spherical_candidates.clip(min_, max_)
# step towards the original inputs
new_source_directions = originals - spherical_candidates
assert new_source_directions.ndim == 2
new_source_directions_norms = ep.norms.l2(flatten(new_source_directions),
axis=-1)
# length if spherical_candidates would be exactly on the sphere
lengths = source_steps * source_norms
# length including correction for numerical deviation from sphere
lengths = lengths + new_source_directions_norms - source_norms
# make sure the step size is positive
lengths = ep.maximum(lengths, 0)
# normalize the length
lengths = lengths / new_source_directions_norms
lengths = atleast_kd(lengths, new_source_directions.ndim)
candidates = spherical_candidates + lengths * new_source_directions
# clip
candidates = candidates.clip(min_, max_)
# restore shape
candidates = candidates.reshape(shape)
spherical_candidates = spherical_candidates.reshape(shape)
return candidates, spherical_candidates