def __call__()

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

• 32 lines of code
• 1 McCabe index (conditional complexity)

    def __call__(self, x, sigmas, train=True):
        # per image standardization
        N = np.prod(x.shape[1:])
        x = (x - jnp.mean(x, axis=(1, 2, 3), keepdims=True)) / jnp.maximum(
            jnp.std(x, axis=(1, 2, 3), keepdims=True), 1.0 / np.sqrt(N)
        )
        temb = GaussianFourierProjection(embedding_size=128, scale=16)(jnp.log(sigmas))
        temb = nn.Dense(128 * 4)(temb)
        temb = nn.Dense(128 * 4)(nn.swish(temb))

        x = nn.Conv(
            16,
            (3, 3),
            padding="SAME",
            name="init_conv",
            kernel_init=conv_kernel_init_fn,
            use_bias=False,
        )(x)
        x = WideResnetGroup(
            self.blocks_per_group,
            16 * self.channel_multiplier,
            activate_before_residual=True,
        )(x, temb, train)
        x = WideResnetGroup(
            self.blocks_per_group, 32 * self.channel_multiplier, (2, 2)
        )(x, temb, train)
        x = WideResnetGroup(
            self.blocks_per_group, 64 * self.channel_multiplier, (2, 2)
        )(x, temb, train)
        x = activation(x, train=train, name="pre-pool-bn")
        x = nn.avg_pool(x, x.shape[1:3])
        x = x.reshape((x.shape[0], -1))
        x = nn.Dense(self.num_outputs, kernel_init=dense_layer_init_fn)(x)
        return x