def __call__()

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

• 30 lines of code
• 4 McCabe index (conditional complexity)

    def __call__(self, x, temb=None, train=True):
        if self.activate_before_residual:
            x = activation(x, train, name="init_bn")
            orig_x = x
        else:
            orig_x = x

        block_x = x
        if not self.activate_before_residual:
            block_x = activation(block_x, train, name="init_bn")

        block_x = nn.Conv(
            self.channels,
            (3, 3),
            self.strides,
            padding="SAME",
            use_bias=False,
            kernel_init=conv_kernel_init_fn,
            name="conv1",
        )(block_x)

        if temb is not None:
            block_x += nn.Dense(self.channels)(nn.swish(temb))[:, None, None, :]
        block_x = activation(block_x, train=train, name="bn_2")
        block_x = nn.Conv(
            self.channels,
            (3, 3),
            padding="SAME",
            use_bias=False,
            kernel_init=conv_kernel_init_fn,
            name="conv2",
        )(block_x)

        return _output_add(block_x, orig_x)