Coverage for src/flag_gems/ops/tanh.py: 89%

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.utils import pointwise_dynamic, tl_extra_shim 

7 

8pow = tl_extra_shim.pow 

9_tanh = tl_extra_shim.tanh 

10logger = logging.getLogger(__name__) 

11 

12 

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

14@triton.jit 

15def tanh_kernel(x): 

16 return _tanh(x.to(tl.float32)) 

17 

18 

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

20@triton.jit 

21def tanh_backward_kernel(y, dy): 

22 y = y.to(tl.float32) 

23 return dy.to(tl.float32) * (1.0 - y * y) 

24 

25 

26def tanh(self): 

27 logger.debug("GEMS TANH FORWARD") 

28 out = tanh_kernel(self) 

29 return out 

30 

31 

32def tanh_backward(grad_output, output): 

33 logger.debug("GEMS TANH BACKWARD") 

34 in_grad = tanh_backward_kernel(output, grad_output) 

35 return in_grad 

36 

37 

38def tanh_(A): 

39 logger.debug("GEMS TANH_ FORWARD") 

40 out = tanh_kernel(A, out0=A) 

41 return out