xref: /aosp_15_r20/external/pytorch/torch/backends/cusparselt/__init__.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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