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

1import logging 

2 

3import torch 

4import triton 

5 

6from ..utils.pointwise_dynamic import pointwise_dynamic 

7 

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

9 

10 

11@pointwise_dynamic(is_tensor=[True, True, False], promotion_methods=[(0, 1, "DEFAULT")]) 

12@triton.jit 

13def rsub_func(x, y, alpha): 

14 return y - x * alpha 

15 

16 

17@pointwise_dynamic( 

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

19) 

20@triton.jit 

21def rsub_func_tensor_scalar(x, y, alpha): 

22 return y - x * alpha 

23 

24 

25@pointwise_dynamic( 

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

27) 

28@triton.jit 

29def rsub_func_scalar_tensor(x, y, alpha): 

30 return y - x * alpha 

31 

32 

33def rsub(A, B, *, alpha=1): 

34 logger.debug("GEMS SUB") 

35 if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor): 

36 return rsub_func(A, B, alpha) 

37 elif isinstance(A, torch.Tensor): 

38 return rsub_func_tensor_scalar(A, B, alpha) 

39 elif isinstance(B, torch.Tensor): 

40 return rsub_func_scalar_tensor(A, B, alpha) 

41 else: 

42 # Both scalar 

43 return B - A * alpha