From 6177b1f7e1382c86957bf386df26b745bac71141 Mon Sep 17 00:00:00 2001
From: Johannes Fischer <johannes.fischer@kit.edu>
Date: Mon, 2 Aug 2021 11:12:24 +0200
Subject: [PATCH] Add kl_cat

---
 src/metrics.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/src/metrics.py b/src/metrics.py
index c9f838c..175d574 100644
--- a/src/metrics.py
+++ b/src/metrics.py
@@ -127,6 +127,23 @@ def kl_normal(pm, pv, qm, qv):
     return kl
 
 
+def kl_cat(q, log_q, log_p):
+    """
+    Computes the KL divergence between two categorical distributions
+
+    Args:
+        q: tensor: (batch, dim): Categorical distribution parameters
+        log_q: tensor: (batch, dim): Log of q
+        log_p: tensor: (batch, dim): Log of p
+
+    Return:
+        kl: tensor: (batch,) kl between each sample
+    """
+    element_wise = (q * (log_q - log_p))
+    kl = element_wise.sum(-1)
+    return kl
+
+
 def nanmean(v, *args, inplace=False, **kwargs):
     """
     Calculate mean over not nan entries