Merge 77526df into 876588d

harvardnlp · Feb 28, 2020 · 5260354 · 5260354
2 parents 876588d + 77526df
commit 5260354
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torch_struct/cky.py b/torch_struct/cky.py
@@ -78,7 +78,7 @@ def arr(a, b):
         log_Z = semiring.dot(top, roots)
         return semiring.unconvert(log_Z), (term_use, rules, top, span[1:]), beta
 
-    def marginals(self, scores, lengths=None, _autograd=True):
+    def marginals(self, scores, lengths=None, _autograd=True, _raw=False):
         """
         Compute the marginals of a CFG using CKY.
 
@@ -92,6 +92,8 @@ def marginals(self, scores, lengths=None, _autograd=True):
             spans: bxNxT terms, (bxNTx(NT+S)x(NT+S)) rules, bxNT roots
 
         """
+        assert not _raw, "top k > 1 with CKY is not supported."
+
         terms, rules, roots = scores
         batch, N, T = terms.shape
         _, NT, _, _ = rules.shape