def build_variant()

in src/kernels/utils.py [0:0]

• 22 lines of code
• 7 McCabe index (conditional complexity)

def build_variant() -> str:
    import torch

    if torch.version.cuda is not None:
        cuda_version = parse(torch.version.cuda)
        compute_framework = f"cu{cuda_version.major}{cuda_version.minor}"
    elif torch.version.hip is not None:
        rocm_version = parse(torch.version.hip.split("-")[0])
        compute_framework = f"rocm{rocm_version.major}{rocm_version.minor}"
    elif torch.backends.mps.is_available():
        compute_framework = "metal"
    else:
        raise AssertionError(
            "Torch was not compiled with CUDA, Metal, or ROCm enabled."
        )

    torch_version = parse(torch.__version__)
    cpu = platform.machine()
    os = platform.system().lower()

    if os == "darwin":
        cpu = "aarch64" if cpu == "arm64" else cpu
        return f"torch{torch_version.major}{torch_version.minor}-{compute_framework}-{cpu}-{os}"

    cxxabi = "cxx11" if torch.compiled_with_cxx11_abi() else "cxx98"

    return f"torch{torch_version.major}{torch_version.minor}-{cxxabi}-{compute_framework}-{cpu}-{os}"