Kernel Density Estimation (KDE) implementation in pytorch or tensorflow

Question

I found an implementation of the Kernel density estimation in scikit-learn as:

from sklearn.neighbors import KernelDensity
kde = KernelDensity(bandwidth=1.0, kernel='gaussian')
kde.fit(x[:, None])
logprob = kde.score_samples(x_d[:, None])

The problem is I want to use Automatic differentiation to take the derivative of logprob w.r.t. x, so I need to use pytorch or tensorflow.

Is there any implementation of the KDE in pytorch or tensorflow, so I can use AD afterward? Or, How can I calculate the derivative of logprob w.r.t. x with scikit-learn?

Answer 1

Update 2023

There is a TensorFlow implementation available in zfit. This is a library that offers PDFs and is based on top of TensorFlow.

JAX now has an implementation

Since you mentioned autograd, JAX is capable of both autograd and jit compilation and offers now the gaussian KDE function that also SciPy implements.

Update

Just to mention, there are also converters that may work, i.e. tf2jax and the other way around.