Coverage for src/flag_gems/ops/isfinite.py: 80%

1import logging 

2 

3import torch 

4import triton 

5import triton.language as tl 

6 

7from flag_gems.utils import pointwise_dynamic, tl_extra_shim 

8 

9try: 

10 _isfinited = tl_extra_shim.isfinited 

11 _finitef = tl_extra_shim.finitef 

12except Exception: 

13 pass 

14logger = logging.getLogger(__name__) 

15 

16 

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

18@triton.jit 

19def isfinite_func(x): 

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

21 

22 

23def isfinite( 

24 A: torch.Tensor, 

25) -> torch.Tensor: 

26 logger.debug("GEMS ISFINITE") 

27 if A.is_floating_point(): 

28 return isfinite_func(A) 

29 else: 

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