(self)
| 302 | dtype: jnp.dtype = jnp.float32 |
| 303 | |
| 304 | def setup(self): |
| 305 | resnets = [] |
| 306 | for i in range(self.num_layers): |
| 307 | in_channels = self.in_channels if i == 0 else self.out_channels |
| 308 | |
| 309 | res_block = FlaxResnetBlock2D( |
| 310 | in_channels=in_channels, |
| 311 | out_channels=self.out_channels, |
| 312 | dropout=self.dropout, |
| 313 | groups=self.resnet_groups, |
| 314 | dtype=self.dtype, |
| 315 | ) |
| 316 | resnets.append(res_block) |
| 317 | self.resnets = resnets |
| 318 | |
| 319 | if self.add_downsample: |
| 320 | self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) |
| 321 | |
| 322 | def __call__(self, hidden_states, deterministic=True): |
| 323 | for resnet in self.resnets: |
nothing calls this directly
no test coverage detected