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

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.utils import tl_extra_shim 

7 

8from ..utils.pointwise_dynamic import pointwise_dynamic 

9 

10_atan = tl_extra_shim.atan 

11logger = logging.getLogger("flag_gems").getChild(__name__.lstrip(".")) 

12 

13 

14@pointwise_dynamic(promotion_methods=[(0, "INT_TO_FLOAT")]) 

15@triton.jit 

16def atan_kernel(x): 

17 return _atan(x.to(tl.float32)) 

18 

19 

20def atan(A): 

21 logger.debug("GEMS ATAN") 

22 out = atan_kernel(A) 

23 return out 

24 

25 

26def atan_(A): 

27 logger.debug("GEMS ATAN_") 

28 atan_kernel(A, out0=A) 

29 return A