def prox(self, x, step_size=None):
shape = x.shape
norms = utils.bnorm(x)
mask = norms > self.alpha
projected = x.clone().detach()
projected[mask] = utils.bdiv(projected[mask], norms[mask])
return self.alpha * projected.view(shape)
def prox(self, x, step_size=None):
"""Projection onto the L2 ball.
Args:
x: torch.Tensor of shape (batchs_size, *)
tensor to project
step_size: Any
Not used here
Returns:
p: torch.Tensor, same shape as x
projection of x onto the L2 ball.
"""
norms = utils.bnorm(x)
mask = norms > self.alpha
projected = x.clone().detach()
projected[mask] = self.alpha * utils.bdiv(projected[mask], norms[mask])
return projected
def is_feasible(self, x, rtol=1e-5, atol=1e-7):
cosines = utils.bdot(x, self.directions)
return abs(
cosines) >= utils.bnorm(x) * self.cos_angle * (1. + rtol) + atol