Coverage for src/flag_gems/runtime/backend/_sunrise/ops/isnan.py: 0%

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.utils import pointwise_dynamic, tl_extra_shim 

7from flag_gems.utils.pointwise_dynamic import CodeGenConfig 

8 

9_isnan = tl_extra_shim.isnan 

10 

11logger = logging.getLogger(__name__) 

12 

13MAX_GRID_SIZES = (65535, 65535, 65535) 

14config = CodeGenConfig( 

15 max_tile_size=1024, 

16 max_grid_size=MAX_GRID_SIZES, 

17 max_num_warps_per_cta=32, 

18 prefer_block_pointer=True, 

19 prefer_1d_tile=True, 

20) 

21 

22 

23@pointwise_dynamic(promotion_methods=[(0, "ALWAYS_BOOL")], config=config) 

24@triton.jit 

25def isnan_func(x): 

26 return _isnan(x.to(tl.float32)) 

27 

28 

29def isnan(A): 

30 logger.debug("GEMS ISNAN") 

31 return isnan_func(A)