(
x_real: torch.Tensor,
x_imag: torch.Tensor,
y_real: torch.Tensor,
y_imag: torch.Tensor,
)
| 21 | |
| 22 | |
| 23 | def complex_elementwise_mul( |
| 24 | x_real: torch.Tensor, |
| 25 | x_imag: torch.Tensor, |
| 26 | y_real: torch.Tensor, |
| 27 | y_imag: torch.Tensor, |
| 28 | ) -> List[torch.Tensor]: |
| 29 | assert x_real.shape == x_imag.shape, f"x_real.shape={x_real.shape}, x_imag.shape={x_imag.shape}" |
| 30 | assert y_real.shape == y_imag.shape, f"y_real.shape={y_real.shape}, y_imag.shape={y_imag.shape}" |
| 31 | return [ |
| 32 | x_real * y_real - x_imag * y_imag, |
| 33 | x_real * y_imag + x_imag * y_real |
| 34 | ] |
| 35 | |
| 36 | |
| 37 | def apply_rotary_pos_emb(x, cos, sin, position_ids, dagger=False): |
no outgoing calls
no test coverage detected