Coverage for src/flag_gems/runtime/backend/_kunlunxin/ops/polar.py: 0%

1import torch 

2import triton 

3import triton.language as tl 

4 

5from ..utils.pointwise_dynamic import pointwise_dynamic 

6 

7 

8@pointwise_dynamic( 

9 promotion_methods=[ 

10 ((0, 1), "DEFAULT"), 

11 ((0, 1), "DEFAULT"), 

12 ], 

13 num_outputs=2, 

14) 

15@triton.jit 

16def polar_kernel(abs, angle): 

17 real = abs * tl.cos(angle) 

18 imag = abs * tl.sin(angle) 

19 return real, imag 

20 

21 

22def polar(abs, angle): 

23 output = torch.empty((*abs.shape, 2), dtype=abs.dtype, device=abs.device) 

24 

25 polar_kernel(abs, angle, out0=output[..., 0], out1=output[..., 1]) 

26 

27 return torch.view_as_complex(output)