def __call__()

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

• 21 lines of code
• 6 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))
        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:
            if self.conv_shortcut:
                x = conv3x3(x, out_ch)
            else:
                x = NIN(out_ch)(x)

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