+
+class BSQ(nn.Module):
+ def __init__(self, L):
+ super().__init__()
+ self.L = L
+
+ def forward(self, input, indexes=False):
+ norm = input.pow(2).sum(dim=2, keepdim=True).sqrt()
+ u = input / norm
+
+ if indexes:
+ return ((u >= 0).long() * (2 ** torch.arange(self.L))[None, :]).sum(dim=1)
+
+ hat_u = 1 / math.sqrt(self.L) * (2 * (u >= 0).float() - 1)
+ if self.training:
+ self.loss += u.mean(dim=0).tanh().pow(2).mean()
+ return hat_u + u - u.detach()
+ else:
+ return hat_u
+
+
+class RandomBypass(nn.Module):
+ def __init__(self, m, p):
+ super().__init__()
+ self.m = m
+ self.p = p
+
+ def forward(self, x):
+ y = self.m(x)
+
+ if self.training:
+ u = (torch.rand(x.size(0), device=x.device) <= self.p).long()[:, None]
+ return (u * x.flatten(1) + (1 - u) * y.flatten(1)).reshape(x.size())
+ else:
+ return y
+
+
+######################################################################
+