Flax implementation of 2D Upsample layer Args: in_channels (`int`): Input channels dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): Parameters `dtype`
| 59 | |
| 60 | |
| 61 | class FlaxUpsample2D(nn.Module): |
| 62 | """ |
| 63 | Flax implementation of 2D Upsample layer |
| 64 | |
| 65 | Args: |
| 66 | in_channels (`int`): |
| 67 | Input channels |
| 68 | dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): |
| 69 | Parameters `dtype` |
| 70 | """ |
| 71 | |
| 72 | in_channels: int |
| 73 | dtype: jnp.dtype = jnp.float32 |
| 74 | |
| 75 | def setup(self): |
| 76 | self.conv = nn.Conv( |
| 77 | self.in_channels, |
| 78 | kernel_size=(3, 3), |
| 79 | strides=(1, 1), |
| 80 | padding=((1, 1), (1, 1)), |
| 81 | dtype=self.dtype, |
| 82 | ) |
| 83 | |
| 84 | def __call__(self, hidden_states): |
| 85 | batch, height, width, channels = hidden_states.shape |
| 86 | hidden_states = jax.image.resize( |
| 87 | hidden_states, |
| 88 | shape=(batch, height * 2, width * 2, channels), |
| 89 | method="nearest", |
| 90 | ) |
| 91 | hidden_states = self.conv(hidden_states) |
| 92 | return hidden_states |
| 93 | |
| 94 | |
| 95 | class FlaxDownsample2D(nn.Module): |