test_header_only.py

import torch
import torch.utils.cpp_extension as _ce
import tempfile
import os
import subprocess

src = """#include <c10/util/BFloat16.h>
#include <iostream>
int main() {
   std::cout << c10::BFloat16(3.14) << std::endl;
}
"""

compiler = _ce.get_cxx_compiler()
if not _ce.check_compiler_ok_for_platform(compiler):
    print("Ieeee")
with tempfile.TemporaryDirectory() as temp_dir:
    with open(os.path.join(temp_dir, "main.cpp"), "w") as f:
        f.write(src)
    subprocess.check_call([compiler, "-std=c++17"] + [f"-I{i}" for i in _ce.include_paths("cpu")] + ["main.cpp"], cwd=temp_dir)
    subprocess.check_call("./a.out", cwd=temp_dir)