Add max gap for alignment (#41)

Author: srush

torch_struct/alignment.py

@@ -14,10 +14,13 @@
 
 
 class Alignment(_Struct):
-    def __init__(self, semiring=LogSemiring, sparse_rounds=3, local=False):
+    def __init__(
+        self, semiring=LogSemiring, sparse_rounds=3, max_gap=None, local=False
+    ):
         self.semiring = semiring
         self.sparse_rounds = sparse_rounds
         self.local = local
+        self.max_gap = max_gap
 
     def _check_potentials(self, edge, lengths=None):
         batch, N_1, M_1, x = edge.shape
@@ -171,6 +174,9 @@ def pad(v):
 
         for n in range(2, log_MN + 1):
             chart = merge(chart)
+            center = (chart.shape[-1] - 1) // 2
+            if self.max_gap is not None and center > self.max_gap:
+                chart = chart[..., center - self.max_gap : center + self.max_gap + 1]
 
         if self.local:
             v = semiring.sum(semiring.sum(chart[..., 0, Close, Close, Mid, :, :]))

torch_struct/semirings/checkpoint.py

@@ -27,7 +27,6 @@ def forward(ctx, a, b):
 
         @staticmethod
         def backward(ctx, grad_output):
-            print("check", grad_output.shape)
             a, b = ctx.saved_tensors
             with torch.enable_grad():
                 q = cls.matmul(a, b)
@@ -43,7 +42,6 @@ def forward(ctx, a, a_lu, a_ld, b, b_lu, b_ld):
 
         @staticmethod
         def backward(ctx, grad_output):
-            print("check_sparse", grad_output.shape)
             a, b, bands = ctx.saved_tensors
             a_lu, a_ld, b_lu, b_ld = bands.tolist()
             with torch.enable_grad():

torch_struct/test_algorithms.py

@@ -365,6 +365,8 @@ def test_alignment(data):
     print(alpha, count)
     print(mx[0].nonzero())
     # assert torch.isclose(count, alpha).all()
+    struct = model(semiring, max_gap=1)
+    alpha = struct.sum(vals)
 
 
 def test_hmm():