fix semimarkov batching and add tests (#114)

* fix semimarkov batching and add tests

* .

* batched implementation

* add back dp_standard

Author: da03

tests/test_algorithms.py

@@ -519,6 +519,32 @@ def test_hsmm(model_test, semiring):
     partition2 = algorithms[model_test][1].enumerate(semiring, edge)[0]
     # third way: dp using edge scores computed from init/transitions/emission
     partition3 = algorithms[model_test][0](semiring).logpartition(edge)[0]
+    # fourth way: dp_standard using edge scores computed from init/transitions/emission
+    partition4 = algorithms[model_test][0](semiring)._dp_standard(edge)[0]
 
     assert torch.isclose(partition1, partition2).all()
     assert torch.isclose(partition2, partition3).all()
+    assert torch.isclose(partition3, partition4).all()
+
+
+@given(data())
+@pytest.mark.parametrize("model_test", ["SemiMarkov"])
+@pytest.mark.parametrize("semiring", [LogSemiring, MaxSemiring])
+def test_batching_lengths(model_test, semiring, data):
+    "Test batching"
+    gen = Gen(model_test, data, LogSemiring)
+    model, vals, N, batch = gen.model, gen.vals, gen.N, gen.batch
+    lengths = torch.tensor(
+        [data.draw(integers(min_value=2, max_value=N)) for b in range(batch - 1)] + [N]
+    )
+    # first way: batched implementation
+    partition = model(semiring).logpartition(vals, lengths=lengths)[0][0]
+    # second way: unbatched implementation
+    for b in range(batch):
+        vals_b = vals[b:(b + 1), :(lengths[b] - 1)]
+        lengths_b = lengths[b:(b + 1)]
+        partition_b = model(semiring).logpartition(vals_b, lengths=lengths_b)[0][0]
+        assert torch.isclose(partition[b], partition_b).all()
+    # test _dp_standard
+    partition_dp_standard = model(semiring)._dp_standard(vals, lengths=lengths)[0][0]
+    assert torch.isclose(partition, partition_dp_standard).all()

torch_struct/semimarkov.py

@@ -34,7 +34,7 @@ def logpartition(self, log_potentials, lengths=None, force_grad=False):
         )
 
         # Init.
-        mask = torch.zeros(*init.shape).bool()
+        mask = torch.zeros(*init.shape, device=log_potentials.device).bool()
         mask[:, :, :, 0, 0].diagonal(0, -2, -1).fill_(True)
         init = semiring.fill(init, mask, semiring.one)
 
@@ -61,10 +61,13 @@ def logpartition(self, log_potentials, lengths=None, force_grad=False):
         c[:, :, : K - 1, 0] = semiring.sum(
             torch.stack([c.data[:, :, : K - 1, 0], lp[:, :, 1:K]], dim=-1)
         )
-        end = torch.min(lengths) - 1
-        mask = torch.zeros(*init.shape).bool()
+        mask = torch.zeros(*init.shape, device=log_potentials.device).bool()
+        mask_length = torch.arange(bin_N).view(1, bin_N, 1).expand(batch, bin_N, C)
+        mask_length = mask_length.to(log_potentials.device)
         for k in range(1, K - 1):
-            mask[:, :, : end - (k - 1), k - 1, k].diagonal(0, -2, -1).fill_(True)
+            mask_length_k = mask_length < (lengths - 1 - (k - 1)).view(batch, 1, 1)
+            mask_length_k = semiring.convert(mask_length_k)
+            mask[:, :, :, k - 1, k].diagonal(0, -2, -1).masked_fill_(mask_length_k, True)
         init = semiring.fill(init, mask, semiring.one)
 
         K_1 = K - 1
@@ -83,37 +86,37 @@ def logpartition(self, log_potentials, lengths=None, force_grad=False):
         v = semiring.sum(semiring.sum(final[:, :, 0, :, 0, :].contiguous()))
         return v, [log_potentials]
 
-    # def _dp_standard(self, edge, lengths=None, force_grad=False):
-    #     semiring = self.semiring
-    #     ssize = semiring.size()
-    #     edge, batch, N, K, C, lengths = self._check_potentials(edge, lengths)
-    #     edge.requires_grad_(True)
-
-    #     # Init
-    #     # All paths starting at N of len K
-    #     alpha = self._make_chart(1, (batch, N, K, C), edge, force_grad)[0]
-
-    #     # All paths finishing at N with label C
-    #     beta = self._make_chart(N, (batch, C), edge, force_grad)
-    #     semiring.one_(beta[0].data)
-
-    #     # Main.
-    #     for n in range(1, N):
-    #         alpha[:, :, n - 1] = semiring.dot(
-    #             beta[n - 1].view(ssize, batch, 1, 1, C),
-    #             edge[:, :, n - 1].view(ssize, batch, K, C, C),
-    #         )
-
-    #         t = max(n - K, -1)
-    #         f1 = torch.arange(n - 1, t, -1)
-    #         f2 = torch.arange(1, len(f1) + 1)
-    #         beta[n][:] = semiring.sum(
-    #             torch.stack([alpha[:, :, a, b] for a, b in zip(f1, f2)], dim=-1)
-    #         )
-    #     v = semiring.sum(
-    #         torch.stack([beta[l - 1][:, i] for i, l in enumerate(lengths)], dim=1)
-    #     )
-    #     return v, [edge], beta
+    def _dp_standard(self, edge, lengths=None, force_grad=False):
+        semiring = self.semiring
+        ssize = semiring.size()
+        edge, batch, N, K, C, lengths = self._check_potentials(edge, lengths)
+        edge.requires_grad_(True)
+
+        # Init
+        # All paths starting at N of len K
+        alpha = self._make_chart(1, (batch, N, K, C), edge, force_grad)[0]
+
+        # All paths finishing at N with label C
+        beta = self._make_chart(N, (batch, C), edge, force_grad)
+        beta[0] = semiring.fill(beta[0], torch.tensor(True).to(edge.device), semiring.one)
+
+        # Main.
+        for n in range(1, N):
+            alpha[:, :, n - 1] = semiring.dot(
+                beta[n - 1].view(ssize, batch, 1, 1, C),
+                edge[:, :, n - 1].view(ssize, batch, K, C, C),
+            )
+
+            t = max(n - K, -1)
+            f1 = torch.arange(n - 1, t, -1)
+            f2 = torch.arange(1, len(f1) + 1)
+            beta[n][:] = semiring.sum(
+                torch.stack([alpha[:, :, a, b] for a, b in zip(f1, f2)], dim=-1)
+            )
+        v = semiring.sum(
+            torch.stack([beta[l - 1][:, i] for i, l in enumerate(lengths)], dim=1)
+        )
+        return v, [edge], beta
 
     @staticmethod
     def to_parts(sequence, extra, lengths=None):