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

1import logging 

2 

3import torch 

4import triton 

5import triton.language as tl 

6 

7from flag_gems.utils import tl_extra_shim 

8 

9from ..utils.pointwise_dynamic import pointwise_dynamic 

10 

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

12_isfinited = tl_extra_shim.isfinited 

13_finitef = tl_extra_shim.finitef 

14 

15 

16@pointwise_dynamic(is_tensor=[True], promotion_methods=[(0, "ALWAYS_BOOL")]) 

17@triton.jit 

18def isfinite_func(x): 

19 return _isfinited(x) if x.dtype.is_fp64() else _finitef(x.to(tl.float32)) 

20 

21 

22def isfinite( 

23 A: torch.Tensor, 

24) -> torch.Tensor: 

25 logger.debug("GEMS ISFINITE") 

26 if A.is_floating_point(): 

27 return isfinite_func(A) 

28 else: 

29 return torch.full(A.shape, True, dtype=torch.bool, device=A.device)