def __call__()

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

• 21 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:
            h = jax.image.resize(x, (x.shape[0], H * 2, W * 2, C), "nearest")
            if self.with_conv:
                h = conv3x3(h, out_ch)
        else:
            if not self.with_conv:
                h = up_or_down_sampling.upsample_2d(x, self.fir_kernel, factor=2)
            else:
                h = up_or_down_sampling.Conv2d(
                    out_ch,
                    kernel=3,
                    up=True,
                    resample_kernel=self.fir_kernel,
                    use_bias=True,
                    kernel_init=default_init(),
                )(x)

        assert h.shape == (B, 2 * H, 2 * W, out_ch)
        return h