def __call__()

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

• 22 lines of code
• 5 McCabe index (conditional complexity)

    def __call__(self, x):
        B, H, W, C = x.shape
        out_ch = self.out_ch if self.out_ch else C
        if not self.fir:
            if self.with_conv:
                x = conv3x3(x, out_ch, stride=2)
            else:
                x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2), padding="SAME")
        else:
            if not self.with_conv:
                x = up_or_down_sampling.downsample_2d(x, self.fir_kernel, factor=2)
            else:
                x = up_or_down_sampling.Conv2d(
                    out_ch,
                    kernel=3,
                    down=True,
                    resample_kernel=self.fir_kernel,
                    use_bias=True,
                    kernel_init=default_init(),
                )(x)

        assert x.shape == (B, H // 2, W // 2, out_ch)
        return x