Coverage for src/flag_gems/runtime/backend/_cambricon/__init__.py: 11%

1import torch # noqa: F401 

2import torch_mlu # noqa: F401 

3 

4from flag_gems.runtime.backend.backend_utils import VendorInfoBase # noqa: E402 

5 

6from .utils import DEVICE_COUNT # noqa: F401 

7from .utils import MAX_GRID_SIZE_X # noqa: F401 

8from .utils import MAX_GRID_SIZE_Y # noqa: F401 

9from .utils import MAX_GRID_SIZE_Z # noqa: F401 

10from .utils import MAX_GRID_SIZES # noqa: F401 

11from .utils import MAX_NRAM_SIZE # noqa: F401 

12from .utils import TOTAL_CLUSTER_NUM # noqa: F401 

13from .utils import TOTAL_CORE_NUM # noqa: F401 

14 

15try: 

16 from torch_mlu.utils.model_transfer import transfer # noqa: F401 

17except ImportError: 

18 pass 

19 

20vendor_info = VendorInfoBase( 

21 vendor_name="cambricon", 

22 device_name="mlu", 

23 device_query_cmd="cnmon", 

24 dispatch_key="PrivateUse1", 

25) 

26 

27CUSTOMIZED_UNUSED_OPS = ( 

28 "scatter_add_", 

29 "masked_scatter", 

30 "masked_scatter_", 

31) 

32 

33__all__ = ["*"]