1769507702

Build Type

github-actions

Committed by GitHub

Commit Message

<a href="https://github.com/harvardnlp/pytorch-struct/commit/<a class=hub.com/harvardnlp/pytorch-struct/commit/<a class="double-link" href="https://git"><a class=hub.com/harvardnlp/pytorch-struct/commit/<a class="double-link" href="https://git"><a class=hub.com/harvardnlp/pytorch-struct/commit/<a class="double-link" href="https://git"><a class=hub.com/harvardnlp/pytorch-struct/commit/<a class="double-link" href="https://git"><a class=hub.com/harvardnlp/pytorch-struct/commit/26d5a7675e863c4275bc3544659a79f82bd186cd">26d5a7675">&lt;a href=&quot;https://github.com/harvardnlp/pytorch-struct/commit/</a><a class="double-link" href="https://github.com/harvardnlp/pytorch-struct/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/harvardnlp/pytorch-struct/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/harvardnlp/pytorch-struct/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/harvardnlp/pytorch-struct/commit/26d5a7675e863c4275bc3544659a79f82bd186cd">26d5a7675</a><a href="https://github.com/harvardnlp/pytorch-struct/commit/26d5a7675e863c4275bc3544659a79f82bd186cd">&lt;a href=&quot;https://github.com/harvardnlp/pytorch-struct/commit/26d5a7675e863c4275bc3544659a79f82bd186cd&quot;&gt;&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/harvardnlp/pytorch-struct/commit/&amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/harvardnlp/pytorch-struct/commit/&amp;amp;lt;a class=&amp;amp;quot;double-link&amp;amp;quot; href=&amp;amp;quot;https://git&amp;quot;&amp;gt;&amp;amp;lt;a class=&amp;lt;/a&amp;gt;hub.com/harvardnlp/pytorch-struct/commit/26d5a7675e863c4275bc3544659a79f82bd186cd&amp;quot;&amp;gt;26d5a7675&amp;lt;/a&amp;gt;&amp;lt;a href=&amp;quot;https://github.com/harvardnlp/pytorch-struct/commit/26d5a7675e863c4275bc3544659a79f82bd186cd&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/harvardnlp/pytorch-struct/commit/26d5a7675e863c4275bc3544659a79f8&amp;lt;/a&amp;gt;2bd186cd&amp;quot</a>;&gt;&amp;quot;&amp;gt;Merge &amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/harvardnlp/pytorch-struct/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/harvardnlp/pytorch-struct/commit/&lt;a class=&quot;double... (continued)

Pull Request Pull Request #123: add initial pass at arg_constraints and arg validation.

Run Details

15 of 15 new or added lines in 1 file covered. (100.0%)

1161 of 1570 relevant lines covered (73.95%)

0.74 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

80.33

/torch_struct/distributions.py

import torch
from torch.distributions import constraints
from torch.distributions.distribution import Distribution
from torch.distributions.utils import lazy_property
from .linearchain import LinearChain
from .cky import CKY
from .semimarkov import SemiMarkov
from .alignment import Alignment
from .deptree import DepTree, deptree_nonproj, deptree_part
from .cky_crf import CKY_CRF
from .semirings import (
    LogSemiring,
    MaxSemiring,
    EntropySemiring,
    CrossEntropySemiring,
    KLDivergenceSemiring,
    MultiSampledSemiring,
    KMaxSemiring,
    StdSemiring,
    GumbelCRFSemiring,
)


class StructDistribution(Distribution):
    r"""
    Base structured distribution class.

    Dynamic distribution for length N of structures :math:`p(z)`.

    Implemented based on gradient identities from:

    * Inside-outside and forward-backward algorithms are just backprop :cite:`eisner2016inside`
    * Semiring Parsing :cite:`goodman1999semiring`
    * First-and second-order expectation semirings with applications to minimum-risk training on translation forests :cite:`li2009first`

    Parameters:
        log_potentials (tensor, batch_shape x event_shape) :  log-potentials :math:`\phi`
        lengths (long tensor, batch_shape) : integers for length masking
    """
    arg_constraints = {
        "log_potentials": constraints.real,
        "lengths": constraints.nonnegative_integer
    }

    def __init__(self, log_potentials, lengths=None, args={}, validate_args=False):
        batch_shape = log_potentials.shape[:1]
        event_shape = log_potentials.shape[1:]
        self.log_potentials = log_potentials
        self.lengths = lengths
        self.args = args
        super().__init__(batch_shape=batch_shape, event_shape=event_shape, validate_args=validate_args)

    def _new(self, *args, **kwargs):
        return self._param.new(*args, **kwargs)

    def log_prob(self, value):
        """
        Compute log probability over values :math:`p(z)`.

        Parameters:
            value (tensor): One-hot events (*sample_shape x batch_shape x event_shape*)

        Returns:
            log_probs (*sample_shape x batch_shape*)
        """

        d = value.dim()
        batch_dims = range(d - len(self.event_shape))
        v = self._struct().score(
            self.log_potentials,
            value.type_as(self.log_potentials),
            batch_dims=batch_dims,
        )

        return v - self.partition

    @lazy_property
    def entropy(self):
        """
        Compute entropy for distribution :math:`H[z]`.

        Returns:
            entropy (*batch_shape*)
        """

        return self._struct(EntropySemiring).sum(self.log_potentials, self.lengths)

    def cross_entropy(self, other):
        """
        Compute cross-entropy for distribution p(self) and q(other) :math:`H[p, q]`.

        Parameters:
            other : Comparison distribution

        Returns:
            cross entropy (*batch_shape*)
        """

        return self._struct(CrossEntropySemiring).sum(
            [self.log_potentials, other.log_potentials], self.lengths
        )

    def kl(self, other):
        """
        Compute KL-divergence for distribution p(self) and q(other) :math:`KL[p || q] = H[p, q] - H[p]`.

        Parameters:
            other : Comparison distribution

        Returns:
            cross entropy (*batch_shape*)
        """
        return self._struct(KLDivergenceSemiring).sum(
            [self.log_potentials, other.log_potentials], self.lengths
        )

    @lazy_property
    def max(self):
        r"""
        Compute an max for distribution :math:`\max p(z)`.

        Returns:
            max (*batch_shape*)
        """
        return self._struct(MaxSemiring).sum(self.log_potentials, self.lengths)

    @lazy_property
    def argmax(self):
        r"""
        Compute an argmax for distribution :math:`\arg\max p(z)`.

        Returns:
            argmax (*batch_shape x event_shape*)
        """
        return self._struct(MaxSemiring).marginals(self.log_potentials, self.lengths)

    def kmax(self, k):
        r"""
        Compute the k-max for distribution :math:`k\max p(z)`.

        Parameters :
            k : Number of solutions to return

        Returns:
            kmax (*k x batch_shape*)
        """
        with torch.enable_grad():
            return self._struct(KMaxSemiring(k)).sum(
                self.log_potentials, self.lengths, _raw=True
            )

    def topk(self, k):
        r"""
        Compute the k-argmax for distribution :math:`k\max p(z)`.

        Parameters :
            k : Number of solutions to return

        Returns:
            kmax (*k x batch_shape x event_shape*)
        """
        with torch.enable_grad():
            return self._struct(KMaxSemiring(k)).marginals(
                self.log_potentials, self.lengths, _raw=True
            )

    @lazy_property
    def mode(self):
        return self.argmax

    @lazy_property
    def marginals(self):
        """
        Compute marginals for distribution :math:`p(z_t)`.

        Can be used in higher-order calculations, i.e.

        *

        Returns:
            marginals (*batch_shape x event_shape*)
        """
        return self._struct(LogSemiring).marginals(self.log_potentials, self.lengths)

    @lazy_property
    def count(self):
        "Compute the log-partition function."
        ones = torch.ones_like(self.log_potentials)
        ones[self.log_potentials.eq(-float("inf"))] = 0
        return self._struct(StdSemiring).sum(ones, self.lengths)

    def gumbel_crf(self, temperature=1.0):
        with torch.enable_grad():
            st_gumbel = self._struct(GumbelCRFSemiring(temperature)).marginals(
                self.log_potentials, self.lengths
            )
            return st_gumbel

    # @constraints.dependent_property
    # def support(self):
    #     pass

    # @property
    # def param_shape(self):
    #     return self._param.size()

    @lazy_property
    def partition(self):
        "Compute the log-partition function."
        return self._struct(LogSemiring).sum(self.log_potentials, self.lengths)

    def sample(self, sample_shape=torch.Size()):
        r"""
        Compute structured samples from the distribution :math:`z \sim p(z)`.

        Parameters:
            sample_shape (int): number of samples

        Returns:
            samples (*sample_shape x batch_shape x event_shape*)
        """
        assert len(sample_shape) == 1
        nsamples = sample_shape[0]
        samples = []
        for k in range(nsamples):
            if k % 10 == 0:
                sample = self._struct(MultiSampledSemiring).marginals(
                    self.log_potentials, lengths=self.lengths
                )
                sample = sample.detach()
            tmp_sample = MultiSampledSemiring.to_discrete(sample, (k % 10) + 1)
            samples.append(tmp_sample)
        return torch.stack(samples)

    def to_event(self, sequence, extra, lengths=None):
        "Convert simple representation to event."
        return self.struct.to_parts(sequence, extra, lengths=lengths)

    def from_event(self, event):
        "Convert event to simple representation."
        return self.struct.from_parts(event)

    def _struct(self, sr=None):
        return self.struct(sr if sr is not None else LogSemiring)


class LinearChainCRF(StructDistribution):
    r"""
    Represents structured linear-chain CRFs with C classes.

    For reference see:

    * An introduction to conditional random fields :cite:`sutton2012introduction`

    Example application:

    * Bidirectional LSTM-CRF Models for Sequence Tagging :cite:`huang2015bidirectional`


    Event shape is of the form:

    Parameters:
        log_potentials (tensor) : event shape (*(N-1) x C x C*) e.g.
                                  :math:`\phi(n,  z_{n+1}, z_{n})`
        lengths (long tensor) : batch_shape integers for length masking.


    Compact representation: N long tensor in [0, ..., C-1]

    Implementation uses linear-scan, forward-pass only.

    * Parallel Time: :math:`O(\log(N))` parallel merges.
    * Forward Memory: :math:`O(N \log(N) C^2)`

    """

    struct = LinearChain


class AlignmentCRF(StructDistribution):
    r"""
    Represents basic alignment algorithm, i.e. dynamic-time warping, Needleman-Wunsch, and Smith-Waterman.

    Event shape is of the form:

    Parameters:
        log_potentials (tensor) : event_shape (*N x M x 3*), e.g.
                                    :math:`\phi(i, j, op)`
                                  Ops are 0 -> j-1, 1->i-1,j-1, and 2->i-1
        local (bool): if true computes local alignment (Smith-Waterman), else Needleman-Wunsch
        max_gap (int or None): the maximum gap to allow in the dynamic program
        lengths (long tensor) : batch shape integers for length masking.


    Implementation uses convolution and linear-scan. Use max_gap for long sequences.

    * Parallel Time: :math:`O(\log (M + N))` parallel merges.
    * Forward Memory: :math:`O((M+N)^2)`

    """
    struct = Alignment
    arg_constraints = {
        "log_potentials": constraints.real,
        "local": constraints.boolean,
        "max_gap": constraints.nonnegative_integer,
        "lengths": constraints.nonnegative_integer
    }

    def __init__(self, log_potentials, local=False, lengths=None, max_gap=None, validate_args=False):
        self.local = local
        self.max_gap = max_gap
        super().__init__(log_potentials, lengths, validate_args=validate_args)

    def _struct(self, sr=None):
        return self.struct(
            sr if sr is not None else LogSemiring, self.local, max_gap=self.max_gap
        )


class HMM(StructDistribution):
    r"""
    Represents hidden-markov smoothing with C hidden states.

    Event shape is of the form:

    Parameters:
        transition (tensor): log-probabilities (*C X C*) :math:`p(z_n| z_n-1)`
        emission (tensor): log-probabilities (*V x C*)  :math:`p(x_n| z_n)`
        init (tensor): log-probabilities (*C*) :math:`p(z_1)`
        observations (long tensor): indices (*batch x N*) between [0, V-1]

    Compact representation: N long tensor in [0, ..., C-1]

    Implemented as a special case of linear chain CRF.
    """

    def __init__(self, transition, emission, init, observations, lengths=None, validate_args=False):
        log_potentials = HMM.struct.hmm(transition, emission, init, observations)
        super().__init__(log_potentials, lengths, validate_args=validate_args)

    struct = LinearChain


class SemiMarkovCRF(StructDistribution):
    r"""
    Represents a semi-markov or segmental CRF with C classes of max width K

    Event shape is of the form:

    Parameters:
       log_potentials : event shape (*N x K x C x C*) e.g.
                        :math:`\phi(n, k, z_{n+1}, z_{n})`
       lengths (long tensor) : batch shape integers for length masking.

    Compact representation: N long tensor in [-1, 0, ..., C-1]

    Implementation uses linear-scan, forward-pass only.

    * Parallel Time: :math:`O(\log(N))` parallel merges.
    * Forward Memory: :math:`O(N \log(N) C^2 K^2)`

    """

    struct = SemiMarkov


class DependencyCRF(StructDistribution):
    r"""
    Represents a projective dependency CRF.

    Reference:

    * Bilexical grammars and their cubic-time parsing algorithms :cite:`eisner2000bilexical`

    Event shape is of the form:

    Parameters:
       log_potentials (tensor) : event shape (*N x N*) head, child or (*N x N x L*) head,
                                 child, labels with arc scores with root scores on diagonal
                                 e.g. :math:`\phi(i, j)` where :math:`\phi(i, i)` is (root, i).
       lengths (long tensor) : batch shape integers for length masking.


    Compact representation: N long tensor in [0, .. N] (indexing is +1)

    Implementation uses linear-scan, forward-pass only.

    * Parallel Time: :math:`O(N)` parallel merges.
    * Forward Memory: :math:`O(N \log(N) C^2 K^2)`

    """

    def __init__(self, log_potentials, lengths=None, args={}, multiroot=True, validate_args=False):
        super(DependencyCRF, self).__init__(log_potentials, lengths, args, validate_args=validate_args)
        self.struct = DepTree
        setattr(self.struct, "multiroot", multiroot)


class TreeCRF(StructDistribution):
    r"""
    Represents a 0th-order span parser with NT nonterminals. Implemented using a
    fast CKY algorithm.

    For example usage see:

    * A Minimal Span-Based Neural Constituency Parser :cite:`stern2017minimal`

    Event shape is of the form:

    Parameters:
        log_potentials (tensor) : event_shape (*N x N x NT*), e.g.
                                    :math:`\phi(i, j, nt)`
        lengths (long tensor) : batch shape integers for length masking.

    Implementation uses width-batched, forward-pass only

    * Parallel Time: :math:`O(N)` parallel merges.
    * Forward Memory: :math:`O(N^2)`

    Compact representation:  *N x N x NT* long tensor (Same)
    """
    struct = CKY_CRF


class SentCFG(StructDistribution):
    """
    Represents a full generative context-free grammar with
    non-terminals NT and terminals T.

    Event shape is of the form:

    Parameters:
        log_potentials (tuple) : event tuple with event shapes
                         terms (*N x T*)
                         rules (*NT x (NT+T) x (NT+T)*)
                         root  (*NT*)
        lengths (long tensor) : batch shape integers for length masking.

    Implementation uses width-batched, forward-pass only

    * Parallel Time: :math:`O(N)` parallel merges.
    * Forward Memory: :math:`O(N^2 (NT+T))`

    Compact representation:  (*N x N x NT*) long tensor
    """

    struct = CKY

    def __init__(self, log_potentials, lengths=None, validate_args=False):
        batch_shape = log_potentials[0].shape[:1]
        event_shape = log_potentials[0].shape[1:]
        self.log_potentials = log_potentials
        self.lengths = lengths
        super(StructDistribution, self).__init__(
            batch_shape=batch_shape, event_shape=event_shape, validate_args=validate_args
        )


class NonProjectiveDependencyCRF(StructDistribution):
    r"""
    Represents a non-projective dependency CRF.

    For references see:

    * Non-projective dependency parsing using spanning tree algorithms :cite:`mcdonald2005non`
    * Structured prediction models via the matrix-tree theorem :cite:`koo2007structured`

    Event shape is of the form:

    Parameters:
       log_potentials (tensor) : event shape (*N x N*) head, child  with
                                 arc scores with root scores on diagonal e.g.
                                 :math:`\phi(i, j)` where :math:`\phi(i, i)` is (root, i).

    Compact representation: N long tensor in [0, .. N] (indexing is +1)

    Note: Does not currently implement argmax (Chiu-Liu) or sampling.

    """

    arg_constraints = {
        "log_potentials": constraints.real
    }

    def __init__(self, log_potentials, lengths=None, args={}, multiroot=False, validate_args=False):
        super(NonProjectiveDependencyCRF, self).__init__(log_potentials, lengths, args, validate_args=validate_args)
        self.multiroot = multiroot

    @lazy_property
    def marginals(self):
        """
        Compute marginals for distribution :math:`p(z_t)`.

        Algorithm is :math:`O(N^3)` but very fast on batched GPU.

        Returns:
            marginals (*batch_shape x event_shape*)
        """
        return deptree_nonproj(self.log_potentials, self.multiroot, self.lengths)

    def sample(self, sample_shape=torch.Size()):
        raise NotImplementedError()

    @lazy_property
    def partition(self):
        """
        Compute the partition function.
        """
        return deptree_part(self.log_potentials, self.multiroot, self.lengths)

    @lazy_property
    def argmax(self):
        """
        Use Chiu-Liu Algorithm. :math:`O(N^2)`

        (Currently not implemented)
        """
        pass

    @lazy_property
    def entropy(self):
        pass

1	import torch	1✔
2	from torch.distributions import constraints	1✔
3	from torch.distributions.distribution import Distribution	1✔
4	from torch.distributions.utils import lazy_property	1✔
5	from .linearchain import LinearChain	1✔
6	from .cky import CKY	1✔
7	from .semimarkov import SemiMarkov	1✔
8	from .alignment import Alignment	1✔
9	from .deptree import DepTree, deptree_nonproj, deptree_part	1✔
10	from .cky_crf import CKY_CRF	1✔
11	from .semirings import (	1✔
12	LogSemiring,
13	MaxSemiring,
14	EntropySemiring,
15	CrossEntropySemiring,
16	KLDivergenceSemiring,
17	MultiSampledSemiring,
18	KMaxSemiring,
19	StdSemiring,
20	GumbelCRFSemiring,
21	)
22
23
24	class StructDistribution(Distribution):	1✔
25	r"""
26	Base structured distribution class.
27
28	Dynamic distribution for length N of structures :math:`p(z)`.
29
30	Implemented based on gradient identities from:
31
32	* Inside-outside and forward-backward algorithms are just backprop :cite:`eisner2016inside`
33	* Semiring Parsing :cite:`goodman1999semiring`
34	* First-and second-order expectation semirings with applications to minimum-risk training on translation forests :cite:`li2009first`
35
36	Parameters:
37	log_potentials (tensor, batch_shape x event_shape) : log-potentials :math:`\phi`
38	lengths (long tensor, batch_shape) : integers for length masking
39	"""
40	arg_constraints = {	1✔
41	"log_potentials": constraints.real,
42	"lengths": constraints.nonnegative_integer
43	}
44
45	def __init__(self, log_potentials, lengths=None, args={}, validate_args=False):	1✔
46	batch_shape = log_potentials.shape[:1]	1✔
47	event_shape = log_potentials.shape[1:]	1✔
48	self.log_potentials = log_potentials	1✔
49	self.lengths = lengths	1✔
50	self.args = args	1✔
51	super().__init__(batch_shape=batch_shape, event_shape=event_shape, validate_args=validate_args)	1✔
52
53	def _new(self, args, *kwargs):	1✔
54	return self._param.new(args, *kwargs)	×
55
56	def log_prob(self, value):	1✔
57	"""
58	Compute log probability over values :math:`p(z)`.
59
60	Parameters:
61	value (tensor): One-hot events (sample_shape x batch_shape x event_shape)
62
63	Returns:
64	log_probs (sample_shape x batch_shape)
65	"""
66
67	d = value.dim()	1✔
68	batch_dims = range(d - len(self.event_shape))	1✔
69	v = self._struct().score(	1✔
70	self.log_potentials,
71	value.type_as(self.log_potentials),
72	batch_dims=batch_dims,
73	)
74
75	return v - self.partition	1✔
76
77	@lazy_property	1✔
78	def entropy(self):
79	"""
80	Compute entropy for distribution :math:`H[z]`.
81
82	Returns:
83	entropy (batch_shape)
84	"""
85
86	return self._struct(EntropySemiring).sum(self.log_potentials, self.lengths)	1✔
87
88	def cross_entropy(self, other):	1✔
89	"""
90	Compute cross-entropy for distribution p(self) and q(other) :math:`H[p, q]`.
91
92	Parameters:
93	other : Comparison distribution
94
95	Returns:
96	cross entropy (batch_shape)
97	"""
98
99	return self._struct(CrossEntropySemiring).sum(	1✔
100	[self.log_potentials, other.log_potentials], self.lengths
101	)
102
103	def kl(self, other):	1✔
104	"""
105	Compute KL-divergence for distribution p(self) and q(other) :math:`KL[p \|\| q] = H[p, q] - H[p]`.
106
107	Parameters:
108	other : Comparison distribution
109
110	Returns:
111	cross entropy (batch_shape)
112	"""
113	return self._struct(KLDivergenceSemiring).sum(	1✔
114	[self.log_potentials, other.log_potentials], self.lengths
115	)
116
117	@lazy_property	1✔
118	def max(self):
119	r"""
120	Compute an max for distribution :math:`\max p(z)`.
121
122	Returns:
123	max (batch_shape)
124	"""
125	return self._struct(MaxSemiring).sum(self.log_potentials, self.lengths)	×
126
127	@lazy_property	1✔
128	def argmax(self):
129	r"""
130	Compute an argmax for distribution :math:`\arg\max p(z)`.
131
132	Returns:
133	argmax (batch_shape x event_shape)
134	"""
135	return self._struct(MaxSemiring).marginals(self.log_potentials, self.lengths)	1✔
136
137	def kmax(self, k):	1✔
138	r"""
139	Compute the k-max for distribution :math:`k\max p(z)`.
140
141	Parameters :
142	k : Number of solutions to return
143
144	Returns:
145	kmax (k x batch_shape)
146	"""
147	with torch.enable_grad():	1✔
148	return self._struct(KMaxSemiring(k)).sum(	1✔
149	self.log_potentials, self.lengths, _raw=True
150	)
151
152	def topk(self, k):	1✔
153	r"""
154	Compute the k-argmax for distribution :math:`k\max p(z)`.
155
156	Parameters :
157	k : Number of solutions to return
158
159	Returns:
160	kmax (k x batch_shape x event_shape)
161	"""
162	with torch.enable_grad():	1✔
163	return self._struct(KMaxSemiring(k)).marginals(	1✔
164	self.log_potentials, self.lengths, _raw=True
165	)
166
167	@lazy_property	1✔
168	def mode(self):
169	return self.argmax	×
170
171	@lazy_property	1✔
172	def marginals(self):
173	"""
174	Compute marginals for distribution :math:`p(z_t)`.
175
176	Can be used in higher-order calculations, i.e.
177
178	*
179
180	Returns:
181	marginals (batch_shape x event_shape)
182	"""
183	return self._struct(LogSemiring).marginals(self.log_potentials, self.lengths)	1✔
184
185	@lazy_property	1✔
186	def count(self):
187	"Compute the log-partition function."
188	ones = torch.ones_like(self.log_potentials)	1✔
189	ones[self.log_potentials.eq(-float("inf"))] = 0	1✔
190	return self._struct(StdSemiring).sum(ones, self.lengths)	1✔
191
192	def gumbel_crf(self, temperature=1.0):	1✔
193	with torch.enable_grad():	×
194	st_gumbel = self._struct(GumbelCRFSemiring(temperature)).marginals(	×
195	self.log_potentials, self.lengths
196	)
197	return st_gumbel	×
198
199	# @constraints.dependent_property
200	# def support(self):
201	# pass
202
203	# @property
204	# def param_shape(self):
205	# return self._param.size()
206
207	@lazy_property	1✔
208	def partition(self):
209	"Compute the log-partition function."
210	return self._struct(LogSemiring).sum(self.log_potentials, self.lengths)	1✔
211
212	def sample(self, sample_shape=torch.Size()):	1✔
213	r"""
214	Compute structured samples from the distribution :math:`z \sim p(z)`.
215
216	Parameters:
217	sample_shape (int): number of samples
218
219	Returns:
220	samples (sample_shape x batch_shape x event_shape)
221	"""
222	assert len(sample_shape) == 1	1✔
223	nsamples = sample_shape[0]	1✔
224	samples = []	1✔
225	for k in range(nsamples):	1✔
226	if k % 10 == 0:	1✔
227	sample = self._struct(MultiSampledSemiring).marginals(	1✔
228	self.log_potentials, lengths=self.lengths
229	)
230	sample = sample.detach()	1✔
231	tmp_sample = MultiSampledSemiring.to_discrete(sample, (k % 10) + 1)	1✔
232	samples.append(tmp_sample)	1✔
233	return torch.stack(samples)	1✔
234
235	def to_event(self, sequence, extra, lengths=None):	1✔
236	"Convert simple representation to event."
237	return self.struct.to_parts(sequence, extra, lengths=lengths)	×
238
239	def from_event(self, event):	1✔
240	"Convert event to simple representation."
241	return self.struct.from_parts(event)	×
242
243	def _struct(self, sr=None):	1✔
244	return self.struct(sr if sr is not None else LogSemiring)	1✔
245
246
247	class LinearChainCRF(StructDistribution):	1✔
248	r"""
249	Represents structured linear-chain CRFs with C classes.
250
251	For reference see:
252
253	* An introduction to conditional random fields :cite:`sutton2012introduction`
254
255	Example application:
256
257	* Bidirectional LSTM-CRF Models for Sequence Tagging :cite:`huang2015bidirectional`
258
259
260	Event shape is of the form:
261
262	Parameters:
263	log_potentials (tensor) : event shape ((N-1) x C x C) e.g.
264	:math:`\phi(n, z_{n+1}, z_{n})`
265	lengths (long tensor) : batch_shape integers for length masking.
266
267
268	Compact representation: N long tensor in [0, ..., C-1]
269
270	Implementation uses linear-scan, forward-pass only.
271
272	* Parallel Time: :math:`O(\log(N))` parallel merges.
273	* Forward Memory: :math:`O(N \log(N) C^2)`
274
275	"""
276
277	struct = LinearChain	1✔
278
279
280	class AlignmentCRF(StructDistribution):	1✔
281	r"""
282	Represents basic alignment algorithm, i.e. dynamic-time warping, Needleman-Wunsch, and Smith-Waterman.
283
284	Event shape is of the form:
285
286	Parameters:
287	log_potentials (tensor) : event_shape (N x M x 3), e.g.
288	:math:`\phi(i, j, op)`
289	Ops are 0 -> j-1, 1->i-1,j-1, and 2->i-1
290	local (bool): if true computes local alignment (Smith-Waterman), else Needleman-Wunsch
291	max_gap (int or None): the maximum gap to allow in the dynamic program
292	lengths (long tensor) : batch shape integers for length masking.
293
294
295	Implementation uses convolution and linear-scan. Use max_gap for long sequences.
296
297	* Parallel Time: :math:`O(\log (M + N))` parallel merges.
298	* Forward Memory: :math:`O((M+N)^2)`
299
300	"""
301	struct = Alignment	1✔
302	arg_constraints = {	1✔
303	"log_potentials": constraints.real,
304	"local": constraints.boolean,
305	"max_gap": constraints.nonnegative_integer,
306	"lengths": constraints.nonnegative_integer
307	}
308
309	def __init__(self, log_potentials, local=False, lengths=None, max_gap=None, validate_args=False):	1✔
310	self.local = local	×
311	self.max_gap = max_gap	×
312	super().__init__(log_potentials, lengths, validate_args=validate_args)	×
313
314	def _struct(self, sr=None):	1✔
315	return self.struct(	×
316	sr if sr is not None else LogSemiring, self.local, max_gap=self.max_gap
317	)
318
319
320	class HMM(StructDistribution):	1✔
321	r"""
322	Represents hidden-markov smoothing with C hidden states.
323
324	Event shape is of the form:
325
326	Parameters:
327	transition (tensor): log-probabilities (C X C) :math:`p(z_n\| z_n-1)`
328	emission (tensor): log-probabilities (V x C) :math:`p(x_n\| z_n)`
329	init (tensor): log-probabilities (C) :math:`p(z_1)`
330	observations (long tensor): indices (batch x N) between [0, V-1]
331
332	Compact representation: N long tensor in [0, ..., C-1]
333
334	Implemented as a special case of linear chain CRF.
335	"""
336
337	def __init__(self, transition, emission, init, observations, lengths=None, validate_args=False):	1✔
338	log_potentials = HMM.struct.hmm(transition, emission, init, observations)	×
339	super().__init__(log_potentials, lengths, validate_args=validate_args)	×
340
341	struct = LinearChain	1✔
342
343
344	class SemiMarkovCRF(StructDistribution):	1✔
345	r"""
346	Represents a semi-markov or segmental CRF with C classes of max width K
347
348	Event shape is of the form:
349
350	Parameters:
351	log_potentials : event shape (N x K x C x C) e.g.
352	:math:`\phi(n, k, z_{n+1}, z_{n})`
353	lengths (long tensor) : batch shape integers for length masking.
354
355	Compact representation: N long tensor in [-1, 0, ..., C-1]
356
357	Implementation uses linear-scan, forward-pass only.
358
359	* Parallel Time: :math:`O(\log(N))` parallel merges.
360	* Forward Memory: :math:`O(N \log(N) C^2 K^2)`
361
362	"""
363
364	struct = SemiMarkov	1✔
365
366
367	class DependencyCRF(StructDistribution):	1✔
368	r"""
369	Represents a projective dependency CRF.
370
371	Reference:
372
373	* Bilexical grammars and their cubic-time parsing algorithms :cite:`eisner2000bilexical`
374
375	Event shape is of the form:
376
377	Parameters:
378	log_potentials (tensor) : event shape (N x N) head, child or (N x N x L) head,
379	child, labels with arc scores with root scores on diagonal
380	e.g. :math:`\phi(i, j)` where :math:`\phi(i, i)` is (root, i).
381	lengths (long tensor) : batch shape integers for length masking.
382
383
384	Compact representation: N long tensor in [0, .. N] (indexing is +1)
385
386	Implementation uses linear-scan, forward-pass only.
387
388	* Parallel Time: :math:`O(N)` parallel merges.
389	* Forward Memory: :math:`O(N \log(N) C^2 K^2)`
390
391	"""
392
393	def __init__(self, log_potentials, lengths=None, args={}, multiroot=True, validate_args=False):	1✔
394	super(DependencyCRF, self).__init__(log_potentials, lengths, args, validate_args=validate_args)	×
395	self.struct = DepTree	×
396	setattr(self.struct, "multiroot", multiroot)	×
397
398
399	class TreeCRF(StructDistribution):	1✔
400	r"""
401	Represents a 0th-order span parser with NT nonterminals. Implemented using a
402	fast CKY algorithm.
403
404	For example usage see:
405
406	* A Minimal Span-Based Neural Constituency Parser :cite:`stern2017minimal`
407
408	Event shape is of the form:
409
410	Parameters:
411	log_potentials (tensor) : event_shape (N x N x NT), e.g.
412	:math:`\phi(i, j, nt)`
413	lengths (long tensor) : batch shape integers for length masking.
414
415	Implementation uses width-batched, forward-pass only
416
417	* Parallel Time: :math:`O(N)` parallel merges.
418	* Forward Memory: :math:`O(N^2)`
419
420	Compact representation: N x N x NT long tensor (Same)
421	"""
422	struct = CKY_CRF	1✔
423
424
425	class SentCFG(StructDistribution):	1✔
426	"""
427	Represents a full generative context-free grammar with
428	non-terminals NT and terminals T.
429
430	Event shape is of the form:
431
432	Parameters:
433	log_potentials (tuple) : event tuple with event shapes
434	terms (N x T)
435	rules (NT x (NT+T) x (NT+T))
436	root (NT)
437	lengths (long tensor) : batch shape integers for length masking.
438
439	Implementation uses width-batched, forward-pass only
440
441	* Parallel Time: :math:`O(N)` parallel merges.
442	* Forward Memory: :math:`O(N^2 (NT+T))`
443
444	Compact representation: (N x N x NT) long tensor
445	"""
446
447	struct = CKY	1✔
448
449	def __init__(self, log_potentials, lengths=None, validate_args=False):	1✔
450	batch_shape = log_potentials[0].shape[:1]	1✔
451	event_shape = log_potentials[0].shape[1:]	1✔
452	self.log_potentials = log_potentials	1✔
453	self.lengths = lengths	1✔
454	super(StructDistribution, self).__init__(	1✔
455	batch_shape=batch_shape, event_shape=event_shape, validate_args=validate_args
456	)
457
458
459	class NonProjectiveDependencyCRF(StructDistribution):	1✔
460	r"""
461	Represents a non-projective dependency CRF.
462
463	For references see:
464
465	* Non-projective dependency parsing using spanning tree algorithms :cite:`mcdonald2005non`
466	* Structured prediction models via the matrix-tree theorem :cite:`koo2007structured`
467
468	Event shape is of the form:
469
470	Parameters:
471	log_potentials (tensor) : event shape (N x N) head, child with
472	arc scores with root scores on diagonal e.g.
473	:math:`\phi(i, j)` where :math:`\phi(i, i)` is (root, i).
474
475	Compact representation: N long tensor in [0, .. N] (indexing is +1)
476
477	Note: Does not currently implement argmax (Chiu-Liu) or sampling.
478
479	"""
480
481	arg_constraints = {	1✔
482	"log_potentials": constraints.real
483	}
484
485	def __init__(self, log_potentials, lengths=None, args={}, multiroot=False, validate_args=False):	1✔
486	super(NonProjectiveDependencyCRF, self).__init__(log_potentials, lengths, args, validate_args=validate_args)	×
487	self.multiroot = multiroot	×
488
489	@lazy_property	1✔
490	def marginals(self):
491	"""
492	Compute marginals for distribution :math:`p(z_t)`.
493
494	Algorithm is :math:`O(N^3)` but very fast on batched GPU.
495
496	Returns:
497	marginals (batch_shape x event_shape)
498	"""
499	return deptree_nonproj(self.log_potentials, self.multiroot, self.lengths)	×
500
501	def sample(self, sample_shape=torch.Size()):	1✔
502	raise NotImplementedError()	×
503
504	@lazy_property	1✔
505	def partition(self):
506	"""
507	Compute the partition function.
508	"""
509	return deptree_part(self.log_potentials, self.multiroot, self.lengths)	×
510
511	@lazy_property	1✔
512	def argmax(self):
513	"""
514	Use Chiu-Liu Algorithm. :math:`O(N^2)`
515
516	(Currently not implemented)
517	"""
518	pass	×
519
520	@lazy_property	1✔
521	def entropy(self):
522	pass	×

harvardnlp / pytorch-struct / 1769507702

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous