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

1import logging 

2import math 

3 

4import torch 

5import triton 

6 

7from flag_gems.runtime import torch_device_fn 

8 

9from .ones import ones_kernel 

10 

11logger = logging.getLogger(f'flag_gems.runtime._ascend.ops.{__name__.split(".")[-1]}') 

12 

13 

14def ones_like( 

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

16): 

17 logger.debug("GEMS_ASCEND ONES_LIKE") 

18 if device is None: 

19 device = x.device 

20 if dtype is None: 

21 dtype = x.dtype 

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

23 N = x.numel() 

24 BLOCK_SIZE = triton.next_power_of_2(math.ceil(math.sqrt(N))) 

25 grid_fn = lambda meta: (triton.cdiv(N, BLOCK_SIZE),) 

26 with torch_device_fn.device(x.device): 

27 ones_kernel[grid_fn](out, N, BLOCK_SIZE=BLOCK_SIZE) 

28 return out