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 return tl.where( 

17 x.to(tl.float32) > 0, 

18 scale * input_scale * x.to(tl.float32), 

19 scale * alpha * (tl.exp(x.to(tl.float32) * input_scale) - 1), 

20 ) 

21 

22 

23@pointwise_dynamic( 

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

25) 

26@triton.jit 

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

28 x_fp32 = x.to(tl.float32) 

29 grad_pos = grad_output * scale * input_scale 

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

31 

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

33 

34 

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

36 logger.debug("GEMS ELU") 

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

38 

39 

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

41 logger.debug("GEMS ELU_") 

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

43 

44 

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

46 logger.debug("GEMS ELU BACKWARD") 

47 grad_input = elu_backward_kernel( 

48 grad_output, self_or_result, alpha, scale, input_scale 

49 ) 

50 return grad_input