Coverage for src/flag_gems/ops/pow.py: 91%

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.utils import pointwise_dynamic, tl_extra_shim 

7 

8_pow = tl_extra_shim.pow 

9logger = logging.getLogger(__name__) 

10 

11 

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

13@triton.jit 

14def pow_func(x, exponent): 

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

16 

17 

18def pow_tensor_tensor(A, exponent): 

19 logger.debug("GEMS POW_TENSOR_TENSOR") 

20 return pow_func(A, exponent) 

21 

22 

23def pow_tensor_tensor_(A, exponent): 

24 logger.debug("GEMS POW_TENSOR_TENSOR_") 

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

26 

27 

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

29@triton.jit 

30def pow_func_tensor_scalar(x, exponent): 

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

32 

33 

34def pow_tensor_scalar(A, exponent): 

35 logger.debug("GEMS POW_TENSOR_SCALAR") 

36 return pow_func_tensor_scalar(A, exponent) 

37 

38 

39def pow_tensor_scalar_(A, exponent): 

40 logger.debug("GEMS POW_TENSOR_SCALAR_") 

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

42 

43 

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

45@triton.jit 

46def pow_func_scalar_tensor(x, exponent): 

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

48 

49 

50def pow_scalar(A, exponent): 

51 logger.debug("GEMS POW_SCALAR") 

52 return pow_func_scalar_tensor(A, exponent)