def __call__()

in jcm/models/layerspp.py [0:0]

• 32 lines of code
• 11 McCabe index (conditional complexity)

    def __call__(self, x, temb=None, train=True):
        B, H, W, C = x.shape
        out_ch = self.out_ch if self.out_ch else C
        h = self.act(nn.GroupNorm(num_groups=min(x.shape[-1] // 4, 32))(x))

        if self.up:
            if self.fir:
                h = up_or_down_sampling.upsample_2d(h, self.fir_kernel, factor=2)
                x = up_or_down_sampling.upsample_2d(x, self.fir_kernel, factor=2)
            else:
                h = up_or_down_sampling.naive_upsample_2d(h, factor=2)
                x = up_or_down_sampling.naive_upsample_2d(x, factor=2)
        elif self.down:
            if self.fir:
                h = up_or_down_sampling.downsample_2d(h, self.fir_kernel, factor=2)
                x = up_or_down_sampling.downsample_2d(x, self.fir_kernel, factor=2)
            else:
                h = up_or_down_sampling.naive_downsample_2d(h, factor=2)
                x = up_or_down_sampling.naive_downsample_2d(x, factor=2)

        h = conv3x3(h, out_ch)
        # Add bias to each feature map conditioned on the time embedding
        if temb is not None:
            h += nn.Dense(out_ch, kernel_init=default_init())(self.act(temb))[
                :, None, None, :
            ]

        h = self.act(nn.GroupNorm(num_groups=min(h.shape[-1] // 4, 32))(h))
        h = nn.Dropout(self.dropout)(h, deterministic=not train)
        h = conv3x3(h, out_ch, init_scale=self.init_scale)
        if C != out_ch or self.up or self.down:
            x = conv1x1(x, out_ch)

        if not self.skip_rescale:
            return x + h
        else:
            return (x + h) / np.sqrt(2.0)