(self, pixel_values)
| 151 | |
| 152 | @nn.compact |
| 153 | def __call__(self, pixel_values): |
| 154 | assert pixel_values.shape[1] == pixel_values.shape[2] == self.config.resolution, pixel_values.shape |
| 155 | hidden_states = nn.Conv(self.config.hidden_channels, [3, 3])(pixel_values) |
| 156 | for i_level in range(self.config.num_resolutions): |
| 157 | hidden_states = DownsamplingBlock(self.config, i_level)(hidden_states) |
| 158 | hidden_states = MidBlock( |
| 159 | self.config, self.config.no_attn_mid_block, self.config.dropout |
| 160 | )(hidden_states) |
| 161 | hidden_states = nn.GroupNorm()(hidden_states) |
| 162 | hidden_states = nn.silu(hidden_states) |
| 163 | hidden_states = nn.Conv(self.config.z_channels, [3, 3])(hidden_states) |
| 164 | return hidden_states |
| 165 | |
| 166 | |
| 167 | class Decoder(nn.Module): |
nothing calls this directly
no test coverage detected