def __call__()

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

• 25 lines of code
• 6 McCabe index (conditional complexity)

    def __call__(self, x):
        assert not (self.up and self.down)
        assert self.kernel >= 1 and self.kernel % 2 == 1
        w = get_weight(
            self,
            (self.kernel, self.kernel, x.shape[-1], self.fmaps),
            weight_var=self.weight_var,
            kernel_init=self.kernel_init,
        )
        if self.up:
            x = upsample_conv_2d(x, w, data_format="NHWC", k=self.resample_kernel)
        elif self.down:
            x = conv_downsample_2d(x, w, data_format="NHWC", k=self.resample_kernel)
        else:
            x = jax.lax.conv_general_dilated(
                x,
                w,
                window_strides=(1, 1),
                padding="SAME",
                dimension_numbers=("NHWC", "HWIO", "NHWC"),
            )

        if self.use_bias:
            b = self.param("bias", jnn.initializers.zeros, (x.shape[-1],))
            x = x + b.reshape((1, 1, 1, -1))
        return x