.

harvardnlp · Nov 26, 2019 · 398ba1a · 398ba1a
1 parent 6dfc342
commit 398ba1a
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 2 deletions.
diff --git a/torch_struct/semimarkov.py b/torch_struct/semimarkov.py
@@ -23,6 +23,7 @@ def _dp(self, log_potentials, lengths=None, force_grad=False):
 
         # Setup
         semiring = self.semiring
+        ssize = semiring.size()
         log_potentials.requires_grad_(True)
         log_potentials, batch, N, K, C, lengths = self._check_potentials(
             log_potentials, lengths
@@ -33,9 +34,43 @@ def _dp(self, log_potentials, lengths=None, force_grad=False):
         )
 
         # Init.
+        semiring.one_(init.data[:, :, :, 0, 0].diagonal(0, -2, -1))
+        # for k in range(1, K - 1):
+        #     semiring.one_(init[:, :, : , k - 1, k].diagonal(0, -2, -1))
+
+        # Length mask
+        big = torch.zeros(
+            ssize,
+            batch,
+            bin_N,
+            K,
+            C,
+            C,
+            dtype=log_potentials.dtype,
+            device=log_potentials.device,
+        )
+        big[:, :, : N - 1] = log_potentials
+        c = init[:, :, :].view(ssize, batch * bin_N, K -1, K-1, C, C)
+        lp = big[:, :, :].view(ssize, batch * bin_N, K, C, C)
+        mask = torch.arange(bin_N).view(1, bin_N).expand(batch, bin_N)
+        mask = mask >= (lengths - 1).view(batch, 1)
+        lp.data[:, mask.view(-1)] = semiring.zero
+        c.data[:, (~mask).view(-1)] = semiring.zero
+        c[:, :,  :K-1, 0] = semiring.sum(torch.stack([c[:, :,  :K-1, 0],
+                                                      lp[:, :,  1:K]], dim=-1))
+
+        # for b in range(lengths.shape[0]):
+        #     end = lengths[b] - 1
+        #     for k in range(1, K - 1):
+        #         semiring.one_(init[:, b, : end - (k - 1), k - 1, k].diagonal(0, -2, -1))
+
+        # ks = torch.arange(1, K - 1)
+        for k in range(1, K - 1):
+            semiring.one_(init[:, :, : (k - 1), k - 1, k].diagonal(0, -2, -1))
 
-        init[:, :, :N-1, : (K - 1), 0] = log_potentials[:, :, :N-1, 1:K]
+        # init[:, :, :end, : (K - 1), 0] = log_potentials[:, :, :end, 1:K]
 
+        # init[:, :, :N-1, : (K - 1), 0] = log_potentials[:, :, :N-1, 1:K]
         # for b in range(lengths.shape[0]):
         #     end = lengths[b] - 1
         #     semiring.one_(init[:, b, end:, 0, 0].diagonal(0, 2, 3))

diff --git a/torch_struct/test_algorithms.py b/torch_struct/test_algorithms.py
@@ -142,7 +142,7 @@ def test_cky(data):
 @settings(max_examples=50, deadline=None)
 def test_generic_a(data):
     model = data.draw(
-        sampled_from([Alignment])  # , LinearChain, SemiMarkov, CKY, CKY_CRF, DepTree])
+        sampled_from([SemiMarkov]) #Alignment , LinearChain, SemiMarkov, CKY, CKY_CRF, DepTree])
     )
     semiring = data.draw(sampled_from([LogSemiring, MaxSemiring]))
     struct = model(semiring)