Coverage for src/flag_gems/runtime/backend/_spacemit/ops/sigmoid.py: 0%

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.utils import tl_extra_shim 

7from flag_gems.utils.pointwise_dynamic import pointwise_dynamic 

8 

9logger = logging.getLogger(__name__) 

10exp = tl_extra_shim.exp 

11 

12 

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

14@triton.jit 

15def sigmoid_forward(x): 

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

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

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

19 return 1 / (1 + exp(-x.to(tl.float32))) 

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_SPACEMIT SIGMOID_FORWARD") 

32 output = sigmoid_forward(self) 

33 return output 

34 

35 

36def sigmoid_backward(grad_output, output): 

37 logger.debug("GEMS_SPACEMIT SIGMOID_BACKWARD") 

38 grad_input = sigmoid_backward_kernel(grad_output, output) 

39 return grad_input 

40 

41 

42def sigmoid_(A): 

43 logger.debug("GEMS_SPACEMIT SIGMOID__FORWARD") 

44 out = sigmoid_forward(A, out0=A) 

45 return out