(self)
| 359 | dtype: jnp.dtype = jnp.float32 |
| 360 | |
| 361 | def setup(self): |
| 362 | resnets = [] |
| 363 | for i in range(self.num_layers): |
| 364 | in_channels = self.in_channels if i == 0 else self.out_channels |
| 365 | res_block = FlaxResnetBlock2D( |
| 366 | in_channels=in_channels, |
| 367 | out_channels=self.out_channels, |
| 368 | dropout=self.dropout, |
| 369 | groups=self.resnet_groups, |
| 370 | dtype=self.dtype, |
| 371 | ) |
| 372 | resnets.append(res_block) |
| 373 | |
| 374 | self.resnets = resnets |
| 375 | |
| 376 | if self.add_upsample: |
| 377 | self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) |
| 378 | |
| 379 | def __call__(self, hidden_states, deterministic=True): |
| 380 | for resnet in self.resnets: |
nothing calls this directly
no test coverage detected