def __call__()

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

• 18 lines of code
• 5 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(self.normalize()(x))
        h = ddpm_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(self.normalize()(h))
        h = nn.Dropout(self.dropout)(h, deterministic=not train)
        h = ddpm_conv3x3(h, out_ch, init_scale=0.0)
        if C != out_ch:
            if self.conv_shortcut:
                x = ddpm_conv3x3(x, out_ch)
            else:
                x = NIN(out_ch)(x)
        return x + h