checks

harvardnlp · Apr 8, 2020 · a792af7 · a792af7
1 parent 9f4630d
commit a792af7
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 1 deletion.
diff --git a/torch_struct/distributions.py b/torch_struct/distributions.py
@@ -142,7 +142,7 @@ def marginals(self):
     def count(self):
         "Compute the log-partition function."
         return self._struct(StdSemiring).sum(
-            self.log_potentials.ones_like(self.log_potentials), self.lengths
+            torch.ones_like(self.log_potentials), self.lengths
         )
 
     # @constraints.dependent_property

diff --git a/torch_struct/semirings/semirings.py b/torch_struct/semirings/semirings.py
@@ -3,6 +3,7 @@
 has_genbmm = False
 try:
     import genbmm
+
     has_genbmm = True
 except ImportError:
     pass

diff --git a/torch_struct/test_distributions.py b/torch_struct/test_distributions.py
@@ -46,6 +46,9 @@ def test_simple(data, seed):
     marginals = dist.marginals
     assert ((samples.mean(0) - marginals).abs() < 0.2).all()
 
+    kmax = dist.kmax(5)
+    count = dist.count
+
 
 @given(data(), integers(min_value=1, max_value=20))
 @settings(max_examples=50, deadline=None)