Coverage for src/flag_gems/runtime/backend/_kunlunxin/ops/pow.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 

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

11_pow = tl_extra_shim.pow 

12 

13 

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

15@triton.jit 

16def pow_func(x, exponent): 

17 return _pow(x.to(tl.float32), exponent.to(tl.float32)) 

18 

19 

20def pow_tensor_tensor(A, exponent): 

21 logger.debug("GEMS POW_TENSOR_TENSOR") 

22 return pow_func(A, exponent) 

23 

24 

25def pow_tensor_tensor_(A, exponent): 

26 logger.debug("GEMS POW_TENSOR_TENSOR_") 

27 return pow_func(A, exponent, out0=A) 

28 

29 

30@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "BOOL_TO_LONG")]) 

31@triton.jit 

32def pow_func_tensor_scalar(x, exponent): 

33 return _pow(x.to(tl.float32), exponent.to(tl.float32)) 

34 

35 

36def pow_tensor_scalar(A, exponent): 

37 logger.debug("GEMS POW_TENSOR_SCALAR") 

38 return pow_func_tensor_scalar(A, exponent) 

39 

40 

41def pow_tensor_scalar_(A, exponent): 

42 logger.debug("GEMS POW_TENSOR_SCALAR_") 

43 return pow_func_tensor_scalar(A, exponent, out0=A) 

44 

45 

46@pointwise_dynamic(is_tensor=[False, True], promotion_methods=[(0, 1, "BOOL_TO_LONG")]) 

47@triton.jit 

48def pow_func_scalar_tensor(x, exponent): 

49 return _pow(x.to(tl.float32), exponent.to(tl.float32)) 

50 

51 

52def pow_scalar(A, exponent): 

53 logger.debug("GEMS POW_SCALAR") 

54 return pow_func_scalar_tensor(A, exponent)