Coverage for src/flag_gems/ops/sigmoid.py: 83%

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__) 

9exp2 = tl_extra_shim.exp2 

10 

11 

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

13@triton.jit 

14def sigmoid_forward(x): 

15 # log2e: tl.constexpr = math.log2(math.e) 

16 # triton 3.0.0 disallow calling non-jitted function inside jitted function, even if it is in 

17 # the rhs of an assignment to a constexpr, so we use numeric literal instead to work around this. 

18 log2e: tl.constexpr = 1.4426950408889634 

19 return 1 / (1 + exp2(-x.to(tl.float32) * log2e)) 

20 

21 

22@pointwise_dynamic(promotion_methods=[(0, "INT_TO_FLOAT")]) 

23@triton.jit 

24def sigmoid_backward_kernel(dy, y): 

25 y_f32 = y.to(tl.float32) 

26 dy_f32 = dy.to(tl.float32) 

27 return dy_f32 * (1.0 - y_f32) * y_f32 

28 

29 

30def sigmoid(self): 

31 logger.debug("GEMS SIGMOID FORWARD") 

32 output = sigmoid_forward(self) 

33 return output 

34 

35 

36def sigmoid_backward(grad_output, output): 

37 logger.debug("GEMS SIGMOID BACKWARD") 

38 grad_input = sigmoid_backward_kernel(grad_output, output) 

39 return grad_input 

40 

41 

42def sigmoid_(A): 

43 logger.debug("GEMS SIGMOID_ FORWARD") 

44 out = sigmoid_forward(A, out0=A) 

45 return out