Pytorch/Tensorflow: Compute gradient of Mixture of Gaussians log density
I have a mixture of three Gaussians and would like to compute the gradient of the log-density using Pytorch or Tensorflow. How can I do that?
from numpy import eye, log
from scipy.stats import multivariate_normal as MVN

μs   = [[0, 0], [2, 0], [0, 2]]              # Means 
Σs   = [eye(2), eye(2), eye(2)]              # Covariance Matrices
cs   = [1 / 3] * 3                           # Mixture coefficients
MVNs = [MVN(μ, Σ) for (μ, Σ) in zip(μs, Σs)] # List of Gaussians

log_density = lambda x: log((sum([c * MVN.pdf(x) for (c, MVN) in zip(cs, MVNs)])))

Essentially I would like to compute the gradient of log_density. I tried using autograd.grad but it fails because of the array assignment.
Attempted Pytorch Solution
from torch import tensor, eye, sqrt, zeros, log, exp
from torch.distributions import MultivariateNormal as MVN

μs   = [tensor([0, 0]), tensor([2, 0]), tensor([0, 2])] # Means 
Σs   = [eye(2), eye(2), eye(2)]                         # Covariance Matrices
cs   = [1 / 3] * 3                                      # Mixture coefficients
MVNs = [MVN(μ, Σ) for (μ, Σ) in zip(μs, Σs)]            # List of Gaussians

log_density = lambda x: log((sum([c * exp(MVN.log_prob(x)) for (c, MVN) in zip(cs, MVNs)])))

Attempted Autograd Solution (won't work)
from numpy import eye, log, zeros
from scipy.stats import multivariate_normal as MVN
from autograd import grad

μs   = [[0, 0], [2, 0], [0, 2]]              # Means 
Σs   = [eye(2), eye(2), eye(2)]              # Covariance Matrices
cs   = [1 / 3] * 3                           # Mixture coefficients
MVNs = [MVN(μ, Σ) for (μ, Σ) in zip(μs, Σs)] # List of Gaussians

log_density = lambda x: log((sum([c * MVN.pdf(x) for (c, MVN) in zip(cs, MVNs)])))
gradient = grad(log_density)

# If you try using this gradient function you get an error
gradient(zeros(2))

The error I get is

ValueError: setting an array element with a sequence.

Naive Autograd Solution
There is, of course, a bad Autograd solution that won't scale well. For instance
from autograd.numpy import log, eye, zeros, array
from autograd.scipy.stats import multivariate_normal as MVN
from autograd import grad

μs   = [[0, 0], [2, 0], [0, 2]]              # Means 
Σs   = [eye(2), eye(2), eye(2)]              # Covariance Matrices
cs   = [1 / 3] * 3                           # Mixture coefficients

def log_density(x):
    return log((1/3) * MVN.pdf(x, zeros(2), eye(2)) + (1/3) * MVN.pdf(x, array([2, 0]), eye(2)) + (1/3) * MVN.pdf(x, array([0, 2]), eye(2)))
grad(log_density)(zeros(2))  # Works!


You can do
from torch import tensor, eye, sqrt, zeros, log, exp
from torch.distributions import MultivariateNormal as MVN

μs   = [tensor([0, 0]), tensor([2, 0]), tensor([0, 2])] # Means 
Σs   = [eye(2), eye(2), eye(2)]                         # Covariance Matrices
cs   = [1 / 3] * 3                                      # Mixture coefficients
MVNs = [MVN(μ, Σ) for (μ, Σ) in zip(μs, Σs)]            # List of Gaussians

x = tensor((0.0,0.0), requires_grad=True)
log_density = log((sum([c * exp(MVN.log_prob(x)) for (c, MVN) in zip(cs, MVNs)])))

log_density.backward()
print(x.grad)

which will print the gradient at (0.0,0.0). However as pytorch is not generating a static computation graph, I could not find an easy way to calculate the gradient at another point without rebuilding the computation graph. You could try to use tensorflow, which gives you more control on the computation graphs and allows you to construct a graph for the gradient computation.
Edit With tensorflow you could do something like
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf
import tensorflow_probability as tfp

@tf.function
def mygrad(x):
    print(&quot;building graph&quot;)
    us   = tf.stack([tf.constant([0.0, 0.0]), tf.constant([2., 0.]), tf.constant([0., 2.])])
    covs = tf.stack([tf.eye(2), tf.eye(2), tf.eye(2)])
    cs   = tf.constant([1 / 3] * 3)

    with tf.GradientTape() as gt:
        gt.watch(x)
        log_density = tf.math.log(tf.math.reduce_sum(tfp.distributions.MultivariateNormalTriL(us,covs).prob(x) * cs) )

    return gt.gradient(log_density,x)

print(mygrad(tf.constant([0.0,0.0])).numpy()) #gradient at 0.0,0.0
print(mygrad(tf.constant([1.0,0.0])).numpy()) #gradient at 1.0,0.0

Essentially you do automatic differentiation with the tf.GradientTape and capture the computation graph in a tf.function. There is more background information on the very extensive Tensorflow API documentation.
