class ScalingLayer(Flow): def __init__(self, init_factor=1.0, dim=1): super().__init__() self._scalefactor = Parameter(init_factor * torch.ones(1)) self.dim = dim def _forward(self, x, *cond, **kwargs): n_batch = x.shape[0] y = torch.zeros_like(x) y[:, : self.dim] = x[:, : self.dim] * self._scalefactor y[:, self.dim :] = x[:, self.dim :] return ( y, (self.dim * self._scalefactor.log()).repeat(n_batch, 1), ) def _inverse(self, x, *cond, **kwargs): n_batch = x.shape[0] y = torch.zeros_like(x) y[:, : self.dim] = x[:, : self.dim] / self._scalefactor y[:, self.dim :] = x[:, self.dim :] return ( y, (-self.dim * self._scalefactor.log()).repeat(n_batch, 1), )
class EntropyScalingLayer(Flow): def __init__(self, init_factor=1.0, dim=1): super().__init__() self._scalefactor = Parameter(init_factor * torch.ones(1)) self.dim = dim def _forward(self, x, y, *cond, **kwargs): n_batch = x.shape[0] return ( self._scalefactor * x, y, (self.dim * self._scalefactor.log()).repeat(n_batch, 1), ) def _inverse(self, x, y, *cond, **kwargs): n_batch = x.shape[0] return ( x / self._scalefactor, y, (-self.dim * self._scalefactor.log()).repeat(n_batch, 1), )