Method _forward_list

vtp/models/layers/block.py:235–298

(self, x_list: List[Tensor], rope_list=None, drop_ratio: Optional[float] = None)

        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]
        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:
            branges_scales_1 = [get_branges_scales(x, effective_drop_ratio) for x in x_list]
            indices_1_list = [br for br, _ in branges_scales_1]
            residual_scale_factors = [sc for _, sc in branges_scales_1]
            x_subset_1_list = [x[indices_1] for x, indices_1 in zip(x_list, indices_1_list)]

            if rope_list is not None:
                rope_subset_list = [
                    self._maybe_index_rope(rope, indices_1) for rope, indices_1 in zip(rope_list, indices_1_list)
                ]
            else:
                rope_subset_list = rope_list

            flattened, shapes, num_tokens = cat_keep_shapes(x_subset_1_list)
            norm1 = uncat_with_shapes(self.norm1(flattened), shapes, num_tokens)
            residual_1_list = self.attn.forward_list(norm1, rope_list=rope_subset_list)

            x_attn_list = [
                torch.index_add(
                    x,
                    dim=0,
                    source=self.ls1(residual_1),
                    index=indices_1,
                    alpha=residual_scale_factor,
                )
                for x, residual_1, indices_1, residual_scale_factor in zip(
                    x_list, residual_1_list, indices_1_list, residual_scale_factors
                )
            ]

            branges_scales_2 = [get_branges_scales(x_attn, effective_drop_ratio) for x_attn in x_attn_list]
            indices_2_list = [br for br, _ in branges_scales_2]
            residual_scale_factors_2 = [sc for _, sc in branges_scales_2]
            x_subset_2_list = [x[indices_2] for x, indices_2 in zip(x_attn_list, indices_2_list)]
            flattened, shapes, num_tokens = cat_keep_shapes(x_subset_2_list)
            norm2_flat = self.norm2(flattened)
            norm2_list = uncat_with_shapes(norm2_flat, shapes, num_tokens)

            residual_2_list = self.mlp.forward_list(norm2_list)

            x_ffn = [
                torch.index_add(
                    x_attn,
                    dim=0,
                    source=self.ls2(residual_2),
                    index=indices_2,
                    alpha=residual_scale_factor,
                )
                for x_attn, residual_2, indices_2, residual_scale_factor in zip(
                    x_attn_list, residual_2_list, indices_2_list, residual_scale_factors_2
                )
            ]
        else:
            x_out = []
            for x, rope in zip(x_list, rope_list):