Coverage for src/flag_gems/ops/add.py: 81%

1import logging 

2 

3import torch 

4import triton 

5 

6from flag_gems.utils import pointwise_dynamic 

7 

8logger = logging.getLogger(__name__) 

9 

10 

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

12@triton.jit 

13def add_func(x, y, alpha): 

14 return x + y * alpha 

15 

16 

17@pointwise_dynamic( 

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

19) 

20@triton.jit 

21def add_func_tensor_scalar(x, y, alpha): 

22 return x + y * alpha 

23 

24 

25@pointwise_dynamic( 

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

27) 

28@triton.jit 

29def add_func_scalar_tensor(x, y, alpha): 

30 return x + y * alpha 

31 

32 

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

34 logger.debug("GEMS ADD") 

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

36 if B.device != A.device: 

37 B = B.to(A.device) 

38 return add_func(A, B, alpha) 

39 elif isinstance(A, torch.Tensor): 

40 return add_func_tensor_scalar(A, B, alpha) 

41 elif isinstance(B, torch.Tensor): 

42 return add_func_scalar_tensor(A, B, alpha) 

43 else: 

44 return torch.tensor(A + B * alpha) 

45 

46 

47def add_(A, B, *, alpha=1): 

48 logger.debug("GEMS ADD_") 

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

50 if B.device != A.device: 

51 B = B.to(A.device) 

52 return add_func(A, B, alpha, out0=A) 

53 elif isinstance(A, torch.Tensor): 

54 return add_func_tensor_scalar(A, B, alpha, out0=A) 

55 # elif isinstance(B, torch.Tensor): 

56 # return add_func_scalar_tensor(A, B, alpha, out0=A) 

57 else: 

58 raise ValueError("Unreachable.")