Coverage for src/flag_gems/fused/mhc/__init__.py: 100%

1from flag_gems.fused.mhc.hc_head_fused_kernel import ( 

2 hc_head_fused_kernel, 

3 hc_head_fused_kernel_ref, 

4) 

5from flag_gems.fused.mhc.hc_split_sinkhorn import ( 

6 hc_split_sinkhorn, 

7 mhc_split_sinkhorn_torch_ref, 

8) 

9from flag_gems.fused.mhc.mhc_bwd import mhc_bwd, mhc_bwd_ref, sinkhorn_forward 

10from flag_gems.fused.mhc.mhc_post import mhc_post 

11from flag_gems.fused.mhc.mhc_pre import mhc_pre 

12 

13__all__ = [ 

14 "hc_head_fused_kernel", 

15 "hc_head_fused_kernel_ref", 

16 "hc_split_sinkhorn", 

17 "mhc_bwd", 

18 "mhc_bwd_ref", 

19 "mhc_post", 

20 "mhc_post_ref", 

21 "mhc_pre", 

22 "mhc_pre_ref", 

23 "mhc_split_sinkhorn_torch_ref", 

24 "sinkhorn_forward", 

25]