Coverage for src/flag_gems/runtime/backend/_hygon/ops/pow.py: 0%

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.utils import pointwise_dynamic, tl_extra_shim 

7 

8logger = logging.getLogger(__name__) 

9_pow = tl_extra_shim.pow 

10 

11 

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

13@triton.jit 

14def pow_func(x, exponent): 

15 if x.type.element_ty == tl.bfloat16: 

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

17 elif x.type.element_ty == tl.float16: 

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

19 else: 

20 return _pow(x.to(tl.float64), exponent) 

21 

22 

23def pow_tensor_tensor(A, exponent): 

24 logger.debug("GEMS POW_TENSOR_TENSOR") 

25 return pow_func(A, exponent) 

26 

27 

28def pow_tensor_tensor_(A, exponent): 

29 logger.debug("GEMS POW_TENSOR_TENSOR_") 

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

31 

32 

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

34@triton.jit 

35def pow_func_tensor_scalar(x, exponent): 

36 if x.type.element_ty == tl.bfloat16: 

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

38 elif x.type.element_ty == tl.float16: 

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

40 else: 

41 return _pow(x.to(tl.float64), exponent) 

42 

43 

44def pow_tensor_scalar(A, exponent): 

45 logger.debug("GEMS POW_TENSOR_SCALAR") 

46 return pow_func_tensor_scalar(A, exponent) 

47 

48 

49def pow_tensor_scalar_(A, exponent): 

50 logger.debug("GEMS POW_TENSOR_SCALAR_") 

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

52 

53 

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

55@triton.jit 

56def pow_func_scalar_tensor(x, exponent): 

57 if exponent.type.element_ty == tl.bfloat16: 

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

59 elif exponent.type.element_ty == tl.float16: 

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

61 else: 

62 return _pow(x.to(tl.float64), exponent) 

63 

64 

65def pow_scalar(A, exponent): 

66 logger.debug("GEMS POW_SCALAR") 

67 return pow_func_scalar_tensor(A, exponent)