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

1import logging 

2 

3import torch 

4import triton 

5 

6from flag_gems.ops.ones import ones_kernel 

7from flag_gems.runtime import torch_device_fn 

8 

9logger = logging.getLogger(__name__) 

10 

11 

12def ones_like( 

13 x, *, dtype=None, layout=None, device=None, pin_memory=None, memory_format=None 

14): 

15 logger.debug("GEMS ONES_LIKE") 

16 if device is None: 

17 device = x.device 

18 if dtype is None: 

19 dtype = x.dtype 

20 out = torch.empty_like(x, device=device, dtype=dtype) 

21 N = x.numel() 

22 grid_fn = lambda meta: (triton.cdiv(N, meta["BLOCK_SIZE"]),) 

23 with torch_device_fn.device(x.device): 

24 ones_kernel[grid_fn](out, N, BLOCK_SIZE=1024) 

25 return out