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

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from ..utils.pointwise_dynamic import pointwise_dynamic 

7 

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

9 

10 

11@pointwise_dynamic( 

12 is_tensor=[True, False, False, False], promotion_methods=[(0, "DEFAULT")] 

13) 

14@triton.jit 

15def elu_forward_kernel(x, alpha, scale, input_scale): 

16 x_fp32 = x.to(tl.float32) 

17 return tl.where( 

18 x_fp32 > 0, 

19 scale * input_scale * x_fp32, 

20 scale * alpha * (tl.exp(x_fp32 * input_scale) - 1), 

21 ) 

22 

23 

24@pointwise_dynamic( 

25 is_tensor=[True, True, False, False, False], promotion_methods=[(0, 1, "DEFAULT")] 

26) 

27@triton.jit 

28def elu_backward_kernel(grad_output, x, alpha, scale, input_scale): 

29 x_fp32 = x.to(tl.float32) 

30 grad_pos = grad_output * scale * input_scale 

31 grad_neg = grad_output * scale * alpha * input_scale * tl.exp(x_fp32 * input_scale) 

32 

33 return tl.where(x_fp32 > 0, grad_pos, grad_neg) 

34 

35 

36def elu(A, alpha=1.0, scale=1.0, input_scale=1.0): 

37 logger.debug("GEMS_KUNLUNXIN ELU") 

38 return elu_forward_kernel(A, alpha, scale, input_scale) 

39 

40 

41def elu_(A, alpha=1.0, scale=1.0, input_scale=1.0): 

42 logger.debug("GEMS_KUNLUNXIN ELU_") 

43 return elu_forward_kernel(A, alpha, scale, input_scale, out0=A) 

44 

45 

46def elu_backward(grad_output, alpha, scale, input_scale, is_result, self_or_result): 

47 logger.debug("GEMS_KUNLUNXIN ELU BACKWARD") 

48 grad_input = elu_backward_kernel( 

49 grad_output, self_or_result, alpha, scale, input_scale 

50 ) 

51 return grad_input