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

1import logging 

2 

3import torch 

4import triton 

5 

6from flag_gems.runtime import torch_device_fn 

7 

8from .ones import ones_kernel 

9 

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

11 

12 

13def ones_like( 

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

15): 

16 logger.debug("GEMS ONES_LIKE") 

17 if device is None: 

18 device = x.device 

19 if dtype is None: 

20 dtype = x.dtype 

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

22 N = x.numel() 

23 if N == 0: 

24 return out 

25 grid_fn = (12, 1, 1) 

26 block_size = triton.next_power_of_2(triton.cdiv(N, 12)) 

27 with torch_device_fn.device(x.device): 

28 ones_kernel[grid_fn]( 

29 out, 

30 N, 

31 1.0, 

32 BLOCK_SIZE=block_size, 

33 buffer_size_limit=2048, 

34 isCloseDtypeConvert=True, 

35 ) 

36 return out