Method _forward

vtp/models/layers/block.py:201–233

(self, x: Tensor, rope=None, drop_ratio: Optional[float] = None)

            return sin, cos

    def _forward(self, x: Tensor, rope=None, drop_ratio: Optional[float] = None) -> Tensor:
        b, _, _ = x.shape
        effective_drop_ratio = drop_ratio if drop_ratio is not None else self.sample_drop_ratio
        if self.training and effective_drop_ratio > 0.0:
            indices_1, residual_scale_factor = get_branges_scales(x, effective_drop_ratio)

            x_subset_1 = x[indices_1]
            rope_subset = self._maybe_index_rope(rope, indices_1)
            residual_1 = self.attn(self.norm1(x_subset_1), rope=rope_subset)

            x_attn = torch.index_add(
                x,
                dim=0,
                source=self.ls1(residual_1),
                index=indices_1,
                alpha=residual_scale_factor,
            )
            indices_2, residual_scale_factor = get_branges_scales(x_attn, effective_drop_ratio)
            x_subset_2 = x_attn[indices_2]
            residual_2 = self.mlp(self.norm2(x_subset_2))

            x_ffn = torch.index_add(
                x_attn,
                dim=0,
                source=self.ls2(residual_2),
                index=indices_2,
                alpha=residual_scale_factor,
            )
        else:
            x_attn = x + self.ls1(self.attn(self.norm1(x), rope=rope))
            x_ffn = x_attn + self.ls2(self.mlp(self.norm2(x_attn)))

        return x_ffn

    def _forward_list(self, x_list: List[Tensor], rope_list=None, drop_ratio: Optional[float] = None) -> List[Tensor]:
        b_list = [x.shape[0] for x in x_list]