Hi, I’m trying to define a new sampler based on torch.poisson which is implemented in cpp. The idea is compile it using cpp_extension from torch. I got compilation errors when trying to include Aten module in cpp. The example code and the returned message can be found below. How can I call that cpp function from pytorch?
cpp_source="""

#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/CPUApplyUtils.h>
#include <ATen/Config.h>
#include <ATen/Dispatch.h>
#include <ATen/ExpandUtils.h>
#include <ATen/NativeFunctions.h>
#include <c10/util/Exception.h>
#include <c10/util/math_compat.h>

#include <ATen/Utils.h>
#include <ATen/CPUGenerator.h>
#include <ATen/core/DistributionsHelper.h>
#include <ATen/native/Distributions.h>
#include <ATen/native/DispatchStub.h>
#include <ATen/native/UnaryOps.h>
#include <ATen/NamedTensorUtils.h>

#include <type_traits>
#include <functional>
#include <assert.h>
#include <cpuinfo.h>
#include <float.h>


int64_t sample_poisson(double lambda, at::CPUGenerator* generator) {
  TORCH_CHECK(lambda >= 0, "invalid Poisson rate, expected rate to be non-negative");
  at::uniform_real_distribution<double> standard_uniform(0.0, 1.0);
  
	int64_t X;
	double prod, U, enlam;

	enlam = std::exp(-lambda);
	X = 0;
	prod = 1.0;
	while (1) {
	  U = standard_uniform(generator);
	  prod *= U;
	  if (prod > enlam) {
		X += 1;
	  } else {
		return X;
	  }
	}
}

"""

import torch
from torch.utils.cpp_extension import load_inline
my_module = load_inline(
    name='my_module',
    cpp_sources=[cpp_source],
    functions=['sample_poisson'],
    build_directory='./build',
    verbose=True
)

I got the following error:
Emitting ninja build file ./build.ninja...
Building extension module my_module...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
[1/2] c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=my_module -DTORCH_API_INCLUDE_EXTENSION_H -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/TH -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/THC -isystem /home/eric/anaconda3/envs/zi/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14 -c /home/eric/Desktop/test_1/main.cpp -o main.o 
FAILED: main.o 
c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=my_module -DTORCH_API_INCLUDE_EXTENSION_H -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/TH -isystem /home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/include/THC -isystem /home/eric/anaconda3/envs/zi/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14 -c /home/eric/Desktop/test_1/main.cpp -o main.o 
/home/eric/Desktop/test_1/main.cpp:15:10: fatal error: ATen/CPUGenerator.h: No such file or directory
   15 | #include <ATen/CPUGenerator.h>
      |          ^~~~~~~~~~~~~~~~~~~~~
compilation terminated.
ninja: build stopped: subcommand failed.
Traceback (most recent call last):
  File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1516, in _run_ninja_build
    subprocess.run(
  File "/home/eric/anaconda3/envs/zi/lib/python3.8/subprocess.py", line 512, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.

The above exception was the direct cause of the following exception:
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1119, in load_inline
    return _jit_compile(
  File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1176, in _jit_compile
    _write_ninja_file_and_build_library(
  File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1280, in _write_ninja_file_and_build_library
    _run_ninja_build(
  File "/home/eric/anaconda3/envs/zi/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1538, in _run_ninja_build
    raise RuntimeError(message) from e
RuntimeError: Error building extension 'my_module'