| 304 | |
| 305 | |
| 306 | class Upsample(nn.Module): |
| 307 | with_conv: bool |
| 308 | |
| 309 | @nn.compact |
| 310 | def __call__(self, hidden_states): |
| 311 | B, H, W, C = hidden_states.shape |
| 312 | hidden_states = jax.image.resize( |
| 313 | hidden_states, |
| 314 | (B, H * 2, W * 2, C), |
| 315 | method="nearest" |
| 316 | ) |
| 317 | if self.with_conv: |
| 318 | hidden_states = nn.Conv(hidden_states.shape[-1], [3, 3])(hidden_states) |
| 319 | return hidden_states |
| 320 | |
| 321 | |
| 322 | class UpsamplingBlock(nn.Module): |