How to calculate Normal Distribution over tensors in Burn?

[AspadaX]

In PyTorch, I can use the Normal object. But the Burn Normal enum now only supports f64 inputs. Or I need to manually compute that?

[laggui]

Burn's Tensor::random samples each element independently from a single, global distribution. There is no built-in support for element-wise distributions with different parameters right now.

But there is a way to achieve this. Sharing a relevant snippet from discord that you could adapt:

impl <B: Backend> AnnealingGaussianSampler<B> {

    pub fn logprob(a: Tensor<B,2>, mu: Tensor<B,2>, std: Tensor<B,2>) -> Tensor<B,2> {
        let std2 = std.powf_scalar(2.0);

        let l = (a - mu).powf_scalar(2.0).div(std2.clone() + 1e-6);
        let r = (std2.mul_scalar(2.0*PI)).log();
        (l + r).mul_scalar(-0.5)
    }

    // Returns an action sampled from the distribution, and the logprob of this action being chosen
    // for the provided mean (mu) and learnable stdev parameter
    pub fn forward(&self, mu: Tensor<B,2>) -> (Tensor<B,2>, Tensor<B,2>) {
        // Use Box Muller transform to get gaussian distribution from pairs of uniformly distributed random numbers
        // This is to get a multivariate gauss distribution using tensors of means and stddevs
        // Otherwise, we would need to create N Distribution::Normal(mu,std) for each element

        let uniform1 = Distribution::Uniform(1e-8, 1.0);
        let uniform2 = Distribution::Uniform(0.0, 1.0);
        let r1 = Tensor::random(mu.shape(), uniform1, &mu.device()).detach();
        let r2 = Tensor::random(mu.shape(), uniform2, &mu.device()).mul_scalar(PI*2.0).detach();

        let mag = self.log_stdev.val().exp().mul(r1.log().mul_scalar(-2.0)).sqrt();
        let sample = (mag * r2.cos()) + mu.clone().detach();

        let logprob = AnnealingGaussianSampler::logprob(sample.clone(), mu, self.log_stdev.val().exp());

        (sample, logprob)
    }
}

[AspadaX]

Burn's Tensor::random samples each element independently from a single, global distribution. There is no built-in support for element-wise distributions with different parameters right now.

But there is a way to achieve this. Sharing a relevant snippet from discord that you could adapt:

impl <B: Backend> AnnealingGaussianSampler<B> {

    pub fn logprob(a: Tensor<B,2>, mu: Tensor<B,2>, std: Tensor<B,2>) -> Tensor<B,2> {
        let std2 = std.powf_scalar(2.0);

        let l = (a - mu).powf_scalar(2.0).div(std2.clone() + 1e-6);
        let r = (std2.mul_scalar(2.0*PI)).log();
        (l + r).mul_scalar(-0.5)
    }

    // Returns an action sampled from the distribution, and the logprob of this action being chosen
    // for the provided mean (mu) and learnable stdev parameter
    pub fn forward(&self, mu: Tensor<B,2>) -> (Tensor<B,2>, Tensor<B,2>) {
        // Use Box Muller transform to get gaussian distribution from pairs of uniformly distributed random numbers
        // This is to get a multivariate gauss distribution using tensors of means and stddevs
        // Otherwise, we would need to create N Distribution::Normal(mu,std) for each element

        let uniform1 = Distribution::Uniform(1e-8, 1.0);
        let uniform2 = Distribution::Uniform(0.0, 1.0);
        let r1 = Tensor::random(mu.shape(), uniform1, &mu.device()).detach();
        let r2 = Tensor::random(mu.shape(), uniform2, &mu.device()).mul_scalar(PI*2.0).detach();

        let mag = self.log_stdev.val().exp().mul(r1.log().mul_scalar(-2.0)).sqrt();
        let sample = (mag * r2.cos()) + mu.clone().detach();

        let logprob = AnnealingGaussianSampler::logprob(sample.clone(), mu, self.log_stdev.val().exp());

        (sample, logprob)
    }
}

That really helps. Thank you for sharing!