def project_onto_l1_ball(x: ep.Tensor, eps: ep.Tensor) -> ep.Tensor: """Computes Euclidean projection onto the L1 ball for a batch. [#Duchi08]_ Adapted from the pytorch version by Tony Duan: https://gist.github.com/tonyduan/1329998205d88c566588e57e3e2c0c55 Args: x: Batch of arbitrary-size tensors to project, possibly on GPU eps: radius of l-1 ball to project onto References: ..[#Duchi08] Efficient Projections onto the l1-Ball for Learning in High Dimensions John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra. International Conference on Machine Learning (ICML 2008) """ original_shape = x.shape x = flatten(x) mask = (ep.norms.l1(x, axis=1) <= eps).astype(x.dtype).expand_dims(1) mu = ep.flip(ep.sort(ep.abs(x)), axis=-1).astype(x.dtype) cumsum = ep.cumsum(mu, axis=-1) arange = ep.arange(x, 1, x.shape[1] + 1).astype(x.dtype) rho = (ep.max( ((mu * arange > (cumsum - eps.expand_dims(1)))).astype(x.dtype) * arange, axis=-1, ) - 1) # samples already under norm will have to select rho = ep.maximum(rho, 0) theta = (cumsum[ep.arange(x, x.shape[0]), rho.astype(ep.arange(x, 1).dtype)] - eps) / (rho + 1.0) proj = (ep.abs(x) - theta.expand_dims(1)).clip(min_=0, max_=ep.inf) x = mask * x + (1 - mask) * proj * ep.sign(x) return x.reshape(original_shape)
def _gram_schmidt(self, v: ep.Tensor, ortho_with: ep.Tensor): v_repeated = ep.concatenate([v.expand_dims(0)] * len(ortho_with), axis=0) #inner product gs_coeff = (ortho_with * v_repeated).flatten(1).sum(1) proj = atleast_kd(gs_coeff, ortho_with.ndim) * ortho_with v = v - proj.sum(0) return v / ep.norms.l2(v)
def project_onto_l1_ball(x: ep.Tensor, eps: ep.Tensor): """ Compute Euclidean projection onto the L1 ball for a batch. min ||x - u||_2 s.t. ||u||_1 <= eps Inspired by the corresponding numpy version by Adrien Gaidon. Adapted from the pytorch version by Tony Duan: https://gist.github.com/tonyduan/1329998205d88c566588e57e3e2c0c55 Parameters ---------- x: (batch_size, *) torch array batch of arbitrary-size tensors to project, possibly on GPU eps: float radius of l-1 ball to project onto Returns ------- u: (batch_size, *) torch array batch of projected tensors, reshaped to match the original Notes ----- The complexity of this algorithm is in O(dlogd) as it involves sorting x. References ---------- [1] Efficient Projections onto the l1-Ball for Learning in High Dimensions John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra. International Conference on Machine Learning (ICML 2008) """ original_shape = x.shape x = flatten(x) mask = (ep.norms.l1(x, axis=1) < eps).astype(x.dtype).expand_dims(1) mu = ep.flip(ep.sort(ep.abs(x)), axis=-1) cumsum = ep.cumsum(mu, axis=-1) arange = ep.arange(x, 1, x.shape[1] + 1) rho = ep.max( (mu * arange > (cumsum - eps.expand_dims(1))) * arange, axis=-1) - 1 theta = (cumsum[ep.arange(x, x.shape[0]), rho] - eps) / (rho + 1.0) proj = (ep.abs(x) - theta.expand_dims(1)).clip(min_=0, max_=ep.inf) x = mask * x + (1 - mask) * proj * ep.sign(x) return x.reshape(original_shape)
def test_squeeze(t: Tensor, axis: Optional[AxisAxes]) -> Tensor: t = t.expand_dims(axis=0).expand_dims(axis=1) return ep.squeeze(t, axis=axis)