fix spacing issues

Author: srush

torch_struct/deptree.py

@@ -182,8 +182,8 @@ def deptree_part(arc_scores, multi_root, lengths, eps=1e-5):
         x = x.unsqueeze(2).expand(-1, -1, N)
         mask = torch.transpose(x, 1, 2) * x
         mask = mask.float()
-        mask[mask==0] = float('-inf')
-        mask[mask==1] = 0
+        mask[mask == 0] = float("-inf")
+        mask[mask == 1] = 0
         arc_scores = arc_scores + mask
     input = arc_scores
     eye = torch.eye(input.shape[1], device=input.device)
@@ -194,13 +194,13 @@ def deptree_part(arc_scores, multi_root, lengths, eps=1e-5):
         lap += det_offset
 
     if multi_root:
-        rss = torch.diagonal(input, 0, -2, -1).exp() # root selection scores
+        rss = torch.diagonal(input, 0, -2, -1).exp()  # root selection scores
         lap = lap + torch.diag_embed(rss, offset=0, dim1=-2, dim2=-1)
     else:
         lap[:, 0] = torch.diagonal(input, 0, -2, -1).exp()
     return lap.logdet()
-    
-    
+
+
 def deptree_nonproj(arc_scores, multi_root, lengths, eps=1e-5):
     """
     Compute the marginals of a non-projective dependency tree using the
@@ -228,10 +228,10 @@ def deptree_nonproj(arc_scores, multi_root, lengths, eps=1e-5):
         x = x.unsqueeze(2).expand(-1, -1, N)
         mask = torch.transpose(x, 1, 2) * x
         mask = mask.float()
-        mask[mask==0] = float('-inf')
-        mask[mask==1] = 0
+        mask[mask == 0] = float("-inf")
+        mask[mask == 1] = 0
         arc_scores = arc_scores + mask
-    
+
     input = arc_scores
     eye = torch.eye(input.shape[1], device=input.device)
     laplacian = input.exp() + eps
@@ -241,7 +241,7 @@ def deptree_nonproj(arc_scores, multi_root, lengths, eps=1e-5):
         lap += det_offset
 
     if multi_root:
-        rss = torch.diagonal(input, 0, -2, -1).exp() # root selection scores
+        rss = torch.diagonal(input, 0, -2, -1).exp()  # root selection scores
         lap = lap + torch.diag_embed(rss, offset=0, dim1=-2, dim2=-1)
         inv_laplacian = lap.inverse()
         factor = (
@@ -254,7 +254,9 @@ def deptree_nonproj(arc_scores, multi_root, lengths, eps=1e-5):
         term2 = input.exp().mul(inv_laplacian.transpose(1, 2)).clone()
         output = term1 - term2
         roots_output = (
-            torch.diagonal(input, 0, -2, -1).exp().mul(torch.diagonal(inv_laplacian.transpose(1, 2), 0, -2, -1))
+            torch.diagonal(input, 0, -2, -1)
+            .exp()
+            .mul(torch.diagonal(inv_laplacian.transpose(1, 2), 0, -2, -1))
         )
     else:
         lap[:, 0] = torch.diagonal(input, 0, -2, -1).exp()
@@ -271,7 +273,9 @@ def deptree_nonproj(arc_scores, multi_root, lengths, eps=1e-5):
         term2[:, 0] = 0
         output = term1 - term2
         roots_output = (
-            torch.diagonal(input, 0, -2, -1).exp().mul(inv_laplacian.transpose(1, 2)[:, 0])
+            torch.diagonal(input, 0, -2, -1)
+            .exp()
+            .mul(inv_laplacian.transpose(1, 2)[:, 0])
         )
     output = output + torch.diag_embed(roots_output, 0, -2, -1)
     return output

torch_struct/distributions.py

@@ -466,10 +466,10 @@ class NonProjectiveDependencyCRF(StructDistribution):
     Note: Does not currently implement argmax (Chiu-Liu) or sampling.
 
     """
+
     def __init__(self, log_potentials, lengths=None, args={}, multiroot=False):
         super(NonProjectiveDependencyCRF, self).__init__(log_potentials, lengths, args)
         self.multiroot = multiroot
-        
 
     @lazy_property
     def marginals(self):