Method forward

model/layers/decoder_block.py:104–140

(self, q, kv)

        self.sample_drop_ratio = drop_path

    def forward(self, q, kv):
        def self_attn_residual_func(q: Tensor) -> Tensor:
            return self.ls1(self.self_attn(self.norm1(q)))

        def cross_attn_residual_func(q: Tensor, kv: Tensor) -> Tensor:
            return self.ls2(self.cross_attn(self.q_norm2(q), self.kv_norm2(kv)))

        def ffn_residual_func(q: Tensor) -> Tensor:
            return self.ls3(self.mlp(self.norm3(q)))

        if self.training and self.sample_drop_ratio > 0.1:
            # the overhead is compensated only for a drop path rate larger than 0.1
            q = drop_add_residual_stochastic_depth(
                [q],
                residual_func=self_attn_residual_func,
                sample_drop_ratio=self.sample_drop_ratio,
            )
            q = drop_add_residual_stochastic_depth(
                [q, kv],
                residual_func=cross_attn_residual_func,
                sample_drop_ratio=self.sample_drop_ratio,
            )
            q = drop_add_residual_stochastic_depth(
                [q],
                residual_func=ffn_residual_func,
                sample_drop_ratio=self.sample_drop_ratio,
            )
        elif self.training and self.sample_drop_ratio > 0.0:
            q = q + self.drop_path1(self_attn_residual_func(q))
            q = q + self.drop_path2(cross_attn_residual_func(q, kv))
            q = q + self.drop_path3(ffn_residual_func(q))
        else:
            q = q + self_attn_residual_func(q)
            q = q + cross_attn_residual_func(q, kv)
            q = q + ffn_residual_func(q)

        return q


def drop_add_residual_stochastic_depth(