(self, q, kv)
| 105 | self.sample_drop_ratio = drop_path |
| 106 | |
| 107 | def forward(self, q, kv): |
| 108 | def self_attn_residual_func(q: Tensor) -> Tensor: |
| 109 | return self.ls1(self.self_attn(self.norm1(q))) |
| 110 | |
| 111 | def cross_attn_residual_func(q: Tensor, kv: Tensor) -> Tensor: |
| 112 | return self.ls2(self.cross_attn(self.q_norm2(q), self.kv_norm2(kv))) |
| 113 | |
| 114 | def ffn_residual_func(q: Tensor) -> Tensor: |
| 115 | return self.ls3(self.mlp(self.norm3(q))) |
| 116 | |
| 117 | if self.training and self.sample_drop_ratio > 0.1: |
| 118 | # the overhead is compensated only for a drop path rate larger than 0.1 |
| 119 | q = drop_add_residual_stochastic_depth( |
| 120 | [q], |
| 121 | residual_func=self_attn_residual_func, |
| 122 | sample_drop_ratio=self.sample_drop_ratio, |
| 123 | ) |
| 124 | q = drop_add_residual_stochastic_depth( |
| 125 | [q, kv], |
| 126 | residual_func=cross_attn_residual_func, |
| 127 | sample_drop_ratio=self.sample_drop_ratio, |
| 128 | ) |
| 129 | q = drop_add_residual_stochastic_depth( |
| 130 | [q], |
| 131 | residual_func=ffn_residual_func, |
| 132 | sample_drop_ratio=self.sample_drop_ratio, |
| 133 | ) |
| 134 | elif self.training and self.sample_drop_ratio > 0.0: |
| 135 | q = q + self.drop_path1(self_attn_residual_func(q)) |
| 136 | q = q + self.drop_path2(cross_attn_residual_func(q, kv)) |
| 137 | q = q + self.drop_path3(ffn_residual_func(q)) |
| 138 | else: |
| 139 | q = q + self_attn_residual_func(q) |
| 140 | q = q + cross_attn_residual_func(q, kv) |
| 141 | q = q + ffn_residual_func(q) |
| 142 | |
| 143 | return q |
| 144 | |
| 145 | |
| 146 | def drop_add_residual_stochastic_depth( |
nothing calls this directly
no test coverage detected