(self, x)
| 72 | self.proj_log_var = nn.Linear(flattened_dim, num_latent_dims) |
| 73 | |
| 74 | def __call__(self, x): |
| 75 | x = nn.leaky_relu(self.bn1(self.conv1(x))) |
| 76 | x = nn.leaky_relu(self.bn2(self.conv2(x))) |
| 77 | x = nn.leaky_relu(self.bn3(self.conv3(x))) |
| 78 | x = mx.flatten(x, 1) # flatten all dimensions except batch |
| 79 | |
| 80 | mu = self.proj_mu(x) |
| 81 | logvar = self.proj_log_var(x) |
| 82 | # Ensure this is the std deviation, not variance |
| 83 | sigma = mx.exp(logvar * 0.5) |
| 84 | |
| 85 | # Generate a tensor of random values from a normal distribution |
| 86 | eps = mx.random.normal(sigma.shape) |
| 87 | |
| 88 | # Reparametrization trick to brackpropagate through sampling. |
| 89 | z = eps * sigma + mu |
| 90 | |
| 91 | return z, mu, logvar |
| 92 | |
| 93 | |
| 94 | class Decoder(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected