Implement Transforms by alicanb · Pull Request #4771 · pytorch/pytorch

alicanb · 2018-01-21T23:25:52Z

This implements a class hierarchy of Transforms, a TransformedDistribution class, and LogNormal and HalfNormal classes as examples of TransformedDistributions. Transform objects have .forward() and .inverse() methods. Most transforms are invertible and also have .log_abs_det_jacobian() methods. The forward and inverse methods are bidirectionally memoized so that e.g. when .forward() is called, subsequent calls to both .forward() and .inverse() are free.

Not all transforms are invertible, e.g. AbsTransform is a 2-to-1 mapping, and LogprobTransform is projects an entire dimension away. However both of these have .inverse() implemented as pseudoinverses. cc @apaszke

yf225 · 2018-01-22T00:03:06Z

@pytorchbot retest this please

apaszke

Looks good, but I'd like to discuss a few things before merging this.

Sign in to view

+from torch.distributions.transformed_distribution import TransformedDistribution
+
+
+class HalfNormal(TransformedDistribution):


Sign in to view


+@register_kl(TransformedDistribution, TransformedDistribution)
+def _kl_transformed_transformed(p, q):
+    if p.transforms != q.transforms:


Sign in to view


+@register_kl(LogNormal, LogNormal)
+def _kl_lognormal_lognormal(p, q):
+    return kl_divergence(p.base_dist, q.base_dist)


Sign in to view

+        if isinstance(transforms, Transform):
+            self.transforms = [transforms, ]
+        elif isinstance(transforms, list):
+            for transform in transforms:


Sign in to view

+        try:
+            return self.transforms[-1].codomain
+        except IndexError:
+            return self.base_dist.support


Sign in to view

+    """
+
+    def __init__(self):
+        self._cache = {}


Sign in to view

+    def __init__(self):
+        self._cache = {}
+
+    def forward(self, x):


Sign in to view

+        shape = x.shape
+        for _ in range(self.event_dim):
+            result = result.sum(-1)
+            shape = shape[:-1]


Sign in to view

+class StickBreakingTransform(Transform):
+    """
+    Transform from the simplex to unconstrained of one fewer dimension via a
+    stick-breaking process.


Sign in to view

+    """
+    Transform from the simplex to unconstrained space via `y = log(x)`.
+
+    This is not bijective and cannot be used for HMC. However this acts mostly


Sign in to view

        self.scale = scale
        self.event_dim = event_dim

+    def __eq__(self, other):


alicanb · 2018-01-22T03:49:40Z

I fixed the low-hanging comments. I also realized we forgot to add tests for LogNormal and HalfNormal, so I added them as well. I think scipy's lognorm.entropy() link is incorrect. so I omitted that test (it passes the monte carlo test).

Sign in to view

+        self._cache = {}
+
+    def __eq__(self, other):
+        return type(other) is type(self)


Sign in to view

+        if self.event_dim:
+            # NOTE: no need for contiguous here
+            result = result.view(*result.size()[:-self.event_dim], -1).sum(-1)
+            shape = shape[:-self.event_dim]


apaszke

Looks good. A few minor questions + please revert submodule change.

Sign in to view

+        """
+        return self._call(x)
+
+    def inverse(self, y):


Sign in to view

+    """
+    Inverts a single :class:`Transform`.
+    """
+    __slots__ = ['inv']


Sign in to view

+        return self.inv.inverse(x)
+
+    def _inverse(self, y):
+        return self.inv.__call__(y)


fritzo

Thanks for the detailed review, Adam!

Sign in to view

+        """
+        return self._call(x)
+
+    def inverse(self, y):


Sign in to view

+    """
+    Inverts a single :class:`Transform`.
+    """
+    __slots__ = ['inv']


Sign in to view

+        return self.inv.inverse(x)
+
+    def _inverse(self, y):
+        return self.inv.__call__(y)


apaszke

Looks good, but let's fix reference cycles.

Sign in to view

+        else:
+            raise NotImplementedError('cache_size must be 0 or 1')
+
+    @lazy_property


Sign in to view

+            pass  # default behavior
+        elif cache_size == 1:
+            self._cached_x_y = None, None
+            self.__call__ = self._cached_call


Sign in to view

+
+    @property
+    def inv(self):
+        return self._inv


Sign in to view

+        try:
+            x_old, y_old = self._cached_x_y
+        except AttributeError:
+            return self._call(x)


Sign in to view

+            return y_old
+        y = self._call(x)
+        self._cached_x_y = x, y
+        return y


fritzo · 2018-01-28T19:29:11Z

Thanks for your patience in reviewing, @apaszke!

apaszke · 2018-01-28T19:43:03Z

Happy to help!

apaszke · 2018-01-28T19:50:01Z

BTW should I squash this, or do you guys want to recreate some meaningful history by manual rebasing (squashed commit can only have a single author).

fritzo · 2018-01-28T19:51:26Z

Squash!

apaszke · 2018-01-28T20:19:01Z

Thanks @fritzo @alicanb!!

alicanb and others added 22 commits January 18, 2018 12:25

add bijectors

89bafd6

Merge branch 'master' of github.com:pytorch/pytorch into bijector

292f8c4

add entropy

1af3851

allow chaining of bijectors

ec6d07e

add affinebijector

a90565f

use log_abs_det_jacobian

860488e

Changed to cache-by-default; add x,y args to log_deg_jacobian method

212f56c

Fix TransformedDistribution.rsample

76742c5

Add KL divergence for LogNormal distribution

d9d371b

Add .domain, .codomain constraints to Bijectors; add tests

001967f

Add .support property to TransformedDistribution

bd69f03

Add AffineBijector

e7ea40e

Add SigmoidBijector

3a17d48

change bijectors to transforms, add half-normal

a58378e

Rename-refactor test_distributions.py

3c78ef3

remove test_transformed_distributions

53cff9f

fix halfnormal docs

59ef2e2

remove log_abs_det_jacobian

84bdfd2

add transforms to distributions.__all__

44bb33b

Add Loprob and StickBreaking transforms

c9d0f46

Add docs for transforms

0052da5

Merge branch 'master' into bijector

f24d3a3

apaszke reviewed Jan 22, 2018

View reviewed changes

address low-hanging comments

788b9f6

alicanb commented Jan 22, 2018

View reviewed changes

Comment thread torch/distributions/transforms.py Outdated

self.scale = scale

self.event_dim = event_dim

def __eq__(self, other):

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

apaszke reviewed Jan 22, 2018

View reviewed changes

Comment thread torch/distributions/transforms.py Outdated

if self.event_dim:

# NOTE: no need for contiguous here

result = result.view(*result.size()[:-self.event_dim], -1).sum(-1)

shape = shape[:-self.event_dim]

This comment was marked as off-topic.

Sign in to view

keep python2 happy

b598b73

Fix bugs in lower triangular checks

f40c9fe

apaszke suggested changes Jan 27, 2018

View reviewed changes

fritzo added 3 commits January 28, 2018 00:09

Revert changes to pybind11 submodule

14ea971

Convert line endings to unix format

0cb8624

Address review comments

e5e0192

fritzo reviewed Jan 28, 2018

View reviewed changes

fritzo added 4 commits January 28, 2018 00:53

Add ComposeTransform and identity_transform

fadd822

Add more documentation

e50e9b0

Fix accidental import

54159bc

Fix typos in documentation

6360a3e

apaszke reviewed Jan 28, 2018

View reviewed changes

Avoid @lazy_property in Transform.inv

3e9d143

fritzo reviewed Jan 28, 2018

View reviewed changes

Comment thread torch/distributions/transforms.py Outdated

pass # default behavior

elif cache_size == 1:

self._cached_x_y = None, None

self.__call__ = self._cached_call

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

fritzo added 2 commits January 28, 2018 09:12

Eliminate reference cycle doe to bound methods

acfa889

Fix docs to satisfy new t.inv.inv == t equation

788c044

apaszke reviewed Jan 28, 2018

View reviewed changes

Use weakrefs to achieve t.inv.inv is t

4390873

apaszke reviewed Jan 28, 2018

View reviewed changes

Avoid try-except-AttributeError pattern

2449320

apaszke approved these changes Jan 28, 2018

View reviewed changes

apaszke merged commit 967bceb into pytorch:master Jan 28, 2018

jpchen mentioned this pull request Jan 30, 2018

Create/Reinstate distribution implementations available in PyTorch pyro-ppl/pyro#699

Closed

3 tasks

ezyang added the open source label Jun 24, 2019

laurentdupin pushed a commit to laurentdupin/pytorch that referenced this pull request Apr 24, 2026

Implement Transforms (pytorch#4771)

9ca3611

		from torch.distributions.transformed_distribution import TransformedDistribution


		class HalfNormal(TransformedDistribution):

Conversation

alicanb commented Jan 21, 2018

Uh oh!

yf225 commented Jan 22, 2018

Uh oh!

apaszke left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

alicanb commented Jan 22, 2018

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

apaszke left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

fritzo left a comment