def __call__()

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

• 52 lines of code
• 12 McCabe index (conditional complexity)

    def __call__(self, x, labels, train=True):
        # config parsing
        config = self.config
        act = get_act(config)
        normalize = get_normalization(config)

        nf = config.model.nf
        ch_mult = config.model.ch_mult
        num_res_blocks = config.model.num_res_blocks
        attn_resolutions = config.model.attn_resolutions
        dropout = config.model.dropout
        resamp_with_conv = config.model.resamp_with_conv
        num_resolutions = len(ch_mult)

        AttnBlock = functools.partial(layers.AttnBlock, normalize=normalize)
        ResnetBlock = functools.partial(
            ResnetBlockDDPM, act=act, normalize=normalize, dropout=dropout
        )

        if config.model.conditional:
            # timestep/scale embedding
            timesteps = labels
            temb = layers.get_timestep_embedding(timesteps, nf)
            temb = nn.Dense(nf * 4, kernel_init=default_initializer())(temb)
            temb = nn.Dense(nf * 4, kernel_init=default_initializer())(act(temb))
        else:
            temb = None

        if config.data.centered:
            # Input is in [-1, 1]
            h = x
        else:
            # Input is in [0, 1]
            h = 2 * x - 1.0

        # Downsampling block
        hs = [conv3x3(h, nf)]
        for i_level in range(num_resolutions):
            # Residual blocks for this resolution
            for i_block in range(num_res_blocks):
                h = ResnetBlock(out_ch=nf * ch_mult[i_level])(hs[-1], temb, train)
                if h.shape[1] in attn_resolutions:
                    h = AttnBlock()(h)
                hs.append(h)
            if i_level != num_resolutions - 1:
                hs.append(Downsample(with_conv=resamp_with_conv)(hs[-1]))

        h = hs[-1]
        h = ResnetBlock()(h, temb, train)
        h = AttnBlock()(h)
        h = ResnetBlock()(h, temb, train)

        # Upsampling block
        for i_level in reversed(range(num_resolutions)):
            for i_block in range(num_res_blocks + 1):
                h = ResnetBlock(out_ch=nf * ch_mult[i_level])(
                    jnp.concatenate([h, hs.pop()], axis=-1), temb, train
                )
            if h.shape[1] in attn_resolutions:
                h = AttnBlock()(h)
            if i_level != 0:
                h = Upsample(with_conv=resamp_with_conv)(h)

        assert not hs

        h = act(normalize()(h))
        h = conv3x3(h, x.shape[-1], init_scale=0.0)
        return h