Add KL Divergence helper #7062
Add KL Divergence helper #7062
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #7062 +/- ##
==========================================
+ Coverage 90.17% 90.18% +0.01%
==========================================
Files 101 103 +2
Lines 16932 16952 +20
==========================================
+ Hits 15269 15289 +20
Misses 1663 1663
|
There was a problem hiding this comment.
This looks very good @ferrine! One thing that you definitely need to do before you can merge this is to add a page about this to the documentation. Maybe make a folder for logprob and include a subfolder for the KL divergence, since it will start to grow as more divergences get added.
pymc/logprob/kl_divergence.py
Outdated
| q_inputs: List[TensorVariable], | ||
| p_inputs: List[TensorVariable], | ||
| ): | ||
| _, _, _, q_mu, q_sigma = q_inputs |
There was a problem hiding this comment.
Should probably consider size like moment does?
I found that our |
|
I am also confused about the design choice, since |
|
The base functionality should be in logprob, but the specific implementations should be in
You can but it won't work for things that look like Distributions but are just helpers to create distributions like |
|
Moved kl into a private pymc.distribution._stats because these are functions that will be never used by anyone |
|
I don't like the underscore, why not just |
|
@ricardoV94 can you please reiterate on the review? Did I miss something? |
|
Tests are failing with import issue |
for more information, see https://pre-commit.ci
|
How many pairs of distributions do we expect to actually be able to support? |
Many of them https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/kl_divergence |
| _, _, _, q_mu, q_sigma = q_inputs | ||
| _, _, _, p_mu, p_sigma = p_inputs | ||
| diff_log_scale = pt.log(q_sigma) - pt.log(p_sigma) | ||
| return ( |
There was a problem hiding this comment.
May want to broadcast to size, like we do with moment, if someone does kl_div(pm.Normal.dist(shape=5), pm.Normal.dist(mu=1))
Description
Related Issue
Checklist
Type of change
📚 Documentation preview 📚: https://pymc--7062.org.readthedocs.build/en/7062/