1# mypy: allow-untyped-defs 2from typing import Optional 3 4import torch 5 6 7__all__ = [ 8 "version", 9 "is_available", 10] 11 12try: 13 from torch._C import _cusparselt 14except ImportError: 15 _cusparselt = None # type: ignore[assignment] 16 17__cusparselt_version: Optional[int] = None 18 19if _cusparselt is not None: 20 21 def _init(): 22 global __cusparselt_version 23 if __cusparselt_version is None: 24 __cusparselt_version = _cusparselt.getVersionInt() 25 return True 26 27else: 28 29 def _init(): 30 return False 31 32 33def version() -> Optional[int]: 34 """Return the version of cuSPARSELt""" 35 if not _init(): 36 return None 37 return __cusparselt_version 38 39 40def is_available() -> bool: 41 r"""Return a bool indicating if cuSPARSELt is currently available.""" 42 return torch._C._has_cusparselt 43