Checks for documentation (#91)

Author: srush

requirements.dev.txt

@@ -2,6 +2,7 @@ pytest
 pytest-runner
 hypothesis == 4.38
 flake8
+darglint
 black
 pep8-naming
 dgl

tests/test_distributions.py

@@ -9,15 +9,18 @@
 lint = integers(min_value=2, max_value=10)
 
 
-def enumerate_support(self, expand=True):
+def enumerate_support(dist):
     """
     Compute the full exponential enumeration set.
 
+    Parameters:
+        dist : Distribution
+
     Returns:
         (enum, enum_lengths) - (*tuple cardinality x batch_shape x event_shape*)
     """
-    _, _, edges, enum_lengths = test_lookup[self.struct]().enumerate(
-        self.log_potentials, self.lengths
+    _, _, edges, enum_lengths = test_lookup[dist.struct]().enumerate(
+        dist.log_potentials, dist.lengths
     )
     # if expand:
     #     edges = edges.unsqueeze(1).expand(edges.shape[:1] + self.batch_shape[:1] + edges.shape[1:])

torch_struct/autoregressive.py

@@ -73,6 +73,7 @@ def log_prob(self, value, sparse=False):
 
         Parameters:
             value (tensor): One-hot events (*sample_shape x batch_shape x N*)
+            sparse (bool): implement sparse
 
         Returns:
             log_probs (*sample_shape x batch_shape*)
@@ -207,6 +208,9 @@ def greedy_tempmax(self, alpha):
 
         * Differentiable Scheduled Sampling for Credit Assignment :cite:`goyal2017differentiable`
 
+        Parameters:
+            alpha : alpha param
+
         Returns:
             greedy_path (*batch x N x C*)
             greedy_max (*batch*)
@@ -219,6 +223,9 @@ def beam_topk(self, K):
         """
         Compute "top-k" using beam search
 
+        Parameters:
+            K : top-k
+
         Returns:
             paths (*K x batch x N x C*)
 

torch_struct/cky.py

@@ -86,7 +86,7 @@ def marginals(self, scores, lengths=None, _autograd=True, _raw=False):
             scores : terms : b x n x T
                      rules : b x NT x (NT+T) x (NT+T)
                      root:   b x NT
-            lengths :
+            lengths : lengths in batch
 
         Returns:
             v: b tensor of total sum

torch_struct/deptree.py

@@ -135,6 +135,9 @@ def to_parts(sequence, extra=None, lengths=None):
 
         Parameters:
             sequence : b x N long tensor in [0, N] (indexing is +1)
+            extra : None
+            lengths : lengths of sequences
+
         Returns:
             arcs : b x N x N arc indicators
         """
@@ -156,6 +159,7 @@ def from_parts(arcs):
 
         Parameters:
             arcs : b x N x N arc indicators
+
         Returns:
             sequence : b x N long tensor in [0, N] (indexing is +1)
         """
@@ -212,7 +216,9 @@ def deptree_nonproj(arc_scores, multi_root, lengths, eps=1e-5):
 
     Parameters:
          arc_scores : b x N x N arc scores with root scores on diagonal.
-         semiring
+         multi_root (bool) : multiple roots
+         lengths : length of examples
+         eps (float) : given
 
     Returns:
          arc_marginals : b x N x N.

torch_struct/distributions.py

@@ -500,8 +500,8 @@ def argmax(self):
 
         (Currently not implemented)
         """
-        raise NotImplementedError()
+        pass
 
     @lazy_property
     def entropy(self):
-        raise NotImplementedError()
+        pass

torch_struct/helpers.py

@@ -64,12 +64,12 @@ def _make_chart(self, N, size, potentials, force_grad=False):
             for _ in range(N)
         ]
 
-    def sum(self, edge, lengths=None, _autograd=True, _raw=False):
+    def sum(self, logpotentials, lengths=None, _autograd=True, _raw=False):
         """
         Compute the (semiring) sum over all structures model.
 
         Parameters:
-            params : generic params (see class)
+            logpotentials : generic params (see class)
             lengths: None or b long tensor mask
 
         Returns:
@@ -82,13 +82,13 @@ def sum(self, edge, lengths=None, _autograd=True, _raw=False):
             or not hasattr(self, "_dp_backward")
         ):
 
-            v = self._dp(edge, lengths)[0]
+            v = self._dp(logpotentials, lengths)[0]
             if _raw:
                 return v
             return self.semiring.unconvert(v)
 
         else:
-            v, _, alpha = self._dp(edge, lengths, False)
+            v, _, alpha = self._dp(logpotentials, lengths, False)
 
             class DPManual(Function):
                 @staticmethod
@@ -97,20 +97,23 @@ def forward(ctx, input):
 
                 @staticmethod
                 def backward(ctx, grad_v):
-                    marginals = self._dp_backward(edge, lengths, alpha)
+                    marginals = self._dp_backward(logpotentials, lengths, alpha)
                     return marginals.mul(
                         grad_v.view((grad_v.shape[0],) + tuple([1] * marginals.dim()))
                     )
 
-            return DPManual.apply(edge)
+            return DPManual.apply(logpotentials)
 
-    def marginals(self, edge, lengths=None, _autograd=True, _raw=False, _combine=False):
+    def marginals(
+        self, logpotentials, lengths=None, _autograd=True, _raw=False, _combine=False
+    ):
         """
         Compute the marginals of a structured model.
 
         Parameters:
-            params : generic params (see class)
+            logpotentials : generic params (see class)
             lengths: None or b long tensor mask
+
         Returns:
             marginals: b x (N-1) x C x C table
 
@@ -120,7 +123,7 @@ def marginals(self, edge, lengths=None, _autograd=True, _raw=False, _combine=Fal
             or self.semiring is not LogSemiring
             or not hasattr(self, "_dp_backward")
         ):
-            v, edges, _ = self._dp(edge, lengths=lengths, force_grad=True)
+            v, edges, _ = self._dp(logpotentials, lengths=lengths, force_grad=True)
             if _raw:
                 all_m = []
                 for k in range(v.shape[0]):
@@ -150,8 +153,8 @@ def marginals(self, edge, lengths=None, _autograd=True, _raw=False, _combine=Fal
                 a_m = self._arrange_marginals(marg)
                 return self.semiring.unconvert(a_m)
         else:
-            v, _, alpha = self._dp(edge, lengths=lengths, force_grad=True)
-            return self._dp_backward(edge, lengths, alpha)
+            v, _, alpha = self._dp(logpotentials, lengths=lengths, force_grad=True)
+            return self._dp_backward(logpotentials, lengths, alpha)
 
     @staticmethod
     def to_parts(spans, extra, lengths=None):

torch_struct/linearchain.py

@@ -92,8 +92,9 @@ def to_parts(sequence, extra, lengths=None):
 
         Parameters:
             sequence : b x N long tensor in [0, C-1]
-            C : number of states
+            extra : number of states
             lengths: b long tensor of N values
+
         Returns:
             edge : b x (N-1) x C x C markov indicators
                         (t x z_t x z_{t-1})
@@ -117,6 +118,7 @@ def from_parts(edge):
         Parameters:
             edge : b x (N-1) x C x C markov indicators
                         (t x z_t x z_{t-1})
+
         Returns:
             sequence : b x N long tensor in [0, C-1]
         """

torch_struct/semimarkov.py

@@ -120,8 +120,9 @@ def to_parts(sequence, extra, lengths=None):
 
         Parameters:
             sequence : b x N  long tensors in [-1, 0, C-1]
-            C : number of states
+            extra : number of states
             lengths: b long tensor of N values
+
         Returns:
             edge : b x (N-1) x K x C x C semimarkov potentials
                         (t x z_t x z_{t-1})
@@ -155,6 +156,7 @@ def from_parts(edge):
         Parameters:
             edge : b x (N-1) x K x C x C semimarkov potentials
                     (t x z_t x z_{t-1})
+
         Returns:
             sequence : b x N  long tensors in [-1, 0, C-1]
 

torch_struct/semirings/semirings.py

@@ -148,12 +148,8 @@ def sum(xs, dim=-1):
         return torch.sum(xs, dim=dim)
 
     @classmethod
-    def matmul(cls, a, b, dims=1):
-        """
-        Dot product along last dim.
-
-        (Faster than calling sum and times.)
-        """
+    def matmul(cls, a, b):
+        "Dot product along last dim"
 
         if has_genbmm and isinstance(a, genbmm.BandedMatrix):
             return b.multiply(a.transpose())
@@ -201,11 +197,7 @@ def sparse_sum(xs, dim=-1):
 
 
 def KMaxSemiring(k):
-    """
-    Implements the k-max semiring (kmax, +, [-inf, -inf..], [0, -inf, ...]).
-
-    Gradients give k-argmax.
-    """
+    "Implements the k-max semiring (kmax, +, [-inf, -inf..], [0, -inf, ...])."
 
     class KMaxSemiring(_BaseLog):
         @staticmethod