Is it not possible to autodiff through jnp.histogram2d?

[Peter-Vincent]

I'm trying to train a model where the loss is the KL-divergence between two joint distributions. I'm constructing these densities by using jnp.histogram2d(x_axis,y_axis,bins=bins) on the generated x_axis and y_axis values for the two densities that I'm trying to get to match. However, when I use grad (or value_and_grad) the gradients are always 0. I've tested gradients through other parts of the code and those seem to work, so my final guess is that I can't audodiff through the jnp.histogram2d() function.

Can you confirm this, and perhaps suggest an alternative?

Thanks

[jakevdp]

You can autodiff through histogram2d, but I believe the correct gradient for the counts in each bin is zero. The issue at play is similar to the one discussed in detail in this answer: #8977 (comment)

Regarding how to work around this: I seem to recall there are ways to estimate the KL divergence without binning your data. I can't remember any exact references, but that's where I would start in this case.

[pourion]

A differentiable density map could be coded as a summation of 2d gaussians centered on each point in the 2D/3D domain. You would need exp() and vmap() to write this code which are both autodiff'able by avoiding direct counting or making bins. I don't know if this is helpful in your case though...