Coverage for src/flag_gems/ops/full_like.py: 100%

1import logging 

2 

3import torch 

4 

5from flag_gems.ops.full import check_dtype, full_func, full_func_scalar 

6 

7logger = logging.getLogger(__name__) 

8 

9 

10def full_like( 

11 x, 

12 fill_value, 

13 *, 

14 dtype=None, 

15 layout=None, 

16 device=None, 

17 pin_memory=None, 

18 memory_format=None, 

19): 

20 logger.debug("GEMS FULL_LIKE") 

21 if device is None: 

22 device = x.device 

23 if dtype is None: 

24 dtype = x.dtype 

25 fill_value = check_dtype(fill_value, dtype, device) 

26 size = x.size() 

27 out = torch.empty(size, device=device, dtype=dtype) 

28 if isinstance(fill_value, torch.Tensor): 

29 return full_func(out, fill_value) 

30 else: 

31 return full_func_scalar(out, fill_value)