torch-ext/torch_binding.cpp

#include <torch/extension.h>
#include "torch_binding.h"

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.doc() = "Erdos-Straus solution counting kernel";
  m.def("count", &erdos_straus_count,
        "Count solutions f(p) to 4/p = 1/x + 1/y + 1/z for each prime p",
        py::arg("primes"));
}