Is there an equivalent of torch.distributions.Normal in LibTorch, the C++ API for PyTorch?

Question

I am implementing a policy gradient algorithm with stochastic policies and since "ancillary" non-PyTorch operations are slow in Python, I want to implement the algorithm in C++. Is there a way to implement a normal distribution in the PyTorch C++ API?

Answer 1

The Python implementation actually calls the C++ back-end in the at:: namespace (CPU, CUDA, where I found this). Until the PyTorch team and/or contributors implement a front-end in LibTorch, you can work around it with something like this (I only implemented rsample() and log_prob() because it's what I need for this use case):

constexpr double lz = log(sqrt(2 * M_PI));
class Normal {
    torch::Tensor mean, stddev, var, log_std;
public:
    Normal(const torch::Tensor &mean, const torch::Tensor &std) : mean(mean), stddev(std), var(std * std), log_std(std.log()) {}


    torch::Tensor rsample() {
        auto device = torch::cuda::is_available() ? torch::kCUDA : torch::kCPU;
        auto eps = torch::randn(1).to(device);
        return this->mean + eps * this->stddev;
    }

    torch::Tensor log_prob(const torch::Tensor &value) {
        // log [exp(-(x-mu)^2/(2 sigma^2)) / (sqrt(2 pi) * sigma)] = 
        // = log [exp(-(x-mu)^2/(2 sigma^2))] - log [sqrt(2 pi) * sigma] = 
        // = -(x - mu)^2 / (2 sigma^2) - log(sigma) - log(sqrt(2 pi))
        return -(value - this->mean)*(value - this->mean) / (2 * this->var) - this->log_std - lz;
    }
};