Add kmax function.

harvardnlp · Apr 7, 2020 · 1cfba78 · 1cfba78
1 parent e76b093
commit 1cfba78
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/torch_struct/distributions.py b/torch_struct/distributions.py
@@ -96,9 +96,20 @@ def argmax(self):
         """
         return self._struct(MaxSemiring).marginals(self.log_potentials, self.lengths)
 
-    def topk(self, k):
+    def kmax(self, k):
         r"""
         Compute the k-max for distribution :math:`k\max p(z)`.
+        Returns:
+            kmax (*k x batch_shape*)
+        """
+        with torch.enable_grad():
+            return self._struct(KMaxSemiring(k)).sum(
+                self.log_potentials, self.lengths, _raw=True
+            )
+
+    def topk(self, k):
+        r"""
+        Compute the k-argmax for distribution :math:`k\max p(z)`.
 
         Returns:
             kmax (*k x batch_shape x event_shape*)