(self, c: Config)
| 25 | num_components: int |
| 26 | |
| 27 | def __init__(self, c: Config): |
| 28 | super().__init__() |
| 29 | self.latent_dim = c.latent_dim |
| 30 | self.num_components = c.num_components |
| 31 | self.input_projection = nn.Linear(c.latent_dim, c.dim) |
| 32 | |
| 33 | self.fc_loc = nn.Linear(c.dim, c.num_components * c.latent_dim) |
| 34 | self.fc_scale = nn.Linear(c.dim, c.num_components * c.latent_dim) |
| 35 | self.fc_weight = nn.Linear(c.dim, c.num_components) |
| 36 | |
| 37 | def _square_plus(self, x): |
| 38 | return (x + T.sqrt(T.square(x) + 4)) / 2 |