Fix TransformedDistribution shaping logic #50581

fritzo · 2021-01-15T14:53:50Z

This fixes many bugs involving TransformedDistribution and ComposeTransform when the component transforms changed their event shapes. Part of the fix is to introduce an IndependentTransform analogous to distributions.Independent and constraints.independent, and to introduce methods Transform.forward_shape() and .inverse_shape(). I have followed @fehiepsi's suggestion and replaced .input_event_dim -> .domain.event_dim and .output_event_dim -> .codomain.event_dim. This allows us to deprecate .event_dim as an attribute.

Summary of changes

Fixes TransformDistribution and ComposeTransform shape errors.
Fixes a behavior bug in LogisticNormal.
Fixes kl_divergence(TransformedDistribution, TransformedDistribution)
Adds methods Transform.forward_shape(), .inverse_shape() which are required for correct shape computations in TransformedDistribution and ComposeTransform.
Adds an IndependentTransform.
Adds a ReshapeTransform which is invaluable in testing shape logic in ComposeTransform and TransformedDistribution and which will be used by @stefanwebb flowtorch.
Fixes incorrect default values in constraints.dependent.event_dim.
Documents the .event_dim and .is_discrete attributes.

Changes planned for follow-up PRs

Memoize @constraints.dependent_property as we do with @lazy_property, since we now consult those properties much more often.

Tested

added a test for Dist.support vs Dist(**params).support to ensure static and dynamic attributes agree.
refactoring is covered by existing tests
add test cases for ReshapedTransform
add a test for TransformedDistribution on a wide grid of input shapes
added a regression test for Incorrect shape from torch.distributions.kl.kl_divergence #34859

cc @fehiepsi @feynmanliang @stefanwebb

torch/distributions/transforms.py

torch/distributions/transformed_distribution.py

torch/distributions/transforms.py

neerajprad · 2021-01-22T04:46:23Z

Could you merge in changes from the earlier PR?

torch/distributions/transforms.py

neerajprad · 2021-01-22T19:55:10Z

The changes to the IndependentTransform look great. I'll defer to @stefanwebb and @fehiepsi for the ReshapeTransform related changes. I'm importing this internally to verify that this doesn't break anything on the client side, and will do another closer review later in the evening.

facebook-github-bot

@neerajprad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

codecov · 2021-01-23T00:22:26Z

Codecov Report

Merging #50581 (e8b121e) into master (b2e5617) will decrease coverage by 0.00%.
The diff coverage is 76.61%.

@@            Coverage Diff             @@
##           master   #50581      +/-   ##
==========================================
- Coverage   80.99%   80.99%   -0.01%     
==========================================
  Files        1916     1916              
  Lines      209590   209748     +158     
==========================================
+ Hits       169767   169884     +117     
- Misses      39823    39864      +41

fehiepsi · 2021-01-23T05:17:58Z

torch/distributions/transformed_distribution.py

+        reinterpreted_batch_ndims = domain_event_dim - base_event_dim
+        if reinterpreted_batch_ndims > 0:
+            base_distribution = Independent(base_distribution, reinterpreted_batch_ndims)
+        self.base_dist = base_distribution


pretty clean!

Let me take another pass

fehiepsi

@fritzo You missed the forward_shape and inverse_shape for ExponentTransform and SoftmaxTransform. The former needs broadcasting to include exponent shape. The latter should be like StickBreakingTransform. Other changes look great to me!

neerajprad

Thanks for adding detailed tests! LGTM, pending comments from @fehipsi's.

fritzo · 2021-01-23T13:57:45Z

@fehiepsi @neerajprad thanks for your detailed review!

You missed the forward_shape and inverse_shape for ExponentTransform and SoftmaxTransform

I don't think PyTorch has an ExponentTransform, that must be NumPyro only 😄 . I've added a check to SoftmaxTransform.forward_shape() and .inverse_shape() but note that transform preserves shape exactly, unlike StickBreakingTransform which changes the size of the rightmost axis.

fehiepsi · 2021-01-23T15:29:19Z

Oops... I meant PowerTransform @fritzo...

facebook-github-bot

@neerajprad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

Summary: Fixes #50496 Fixes #34859 Fixes #21596 This fixes many bugs involving `TransformedDistribution` and `ComposeTransform` when the component transforms changed their event shapes. Part of the fix is to introduce an `IndependentTransform` analogous to `distributions.Independent` and `constraints.independent`, and to introduce methods `Transform.forward_shape()` and `.inverse_shape()`. I have followed fehiepsi's suggestion and replaced `.input_event_dim` -> `.domain.event_dim` and `.output_event_dim` -> `.codomain.event_dim`. This allows us to deprecate `.event_dim` as an attribute. ## Summary of changes - Fixes `TransformDistribution` and `ComposeTransform` shape errors. - Fixes a behavior bug in `LogisticNormal`. - Fixes `kl_divergence(TransformedDistribution, TransformedDistribution)` - Adds methods `Transform.forward_shape()`, `.inverse_shape()` which are required for correct shape computations in `TransformedDistribution` and `ComposeTransform`. - Adds an `IndependentTransform`. - Adds a `ReshapeTransform` which is invaluable in testing shape logic in `ComposeTransform` and `TransformedDistribution` and which will be used by stefanwebb flowtorch. - Fixes incorrect default values in `constraints.dependent.event_dim`. - Documents the `.event_dim` and `.is_discrete` attributes. ## Changes planned for follow-up PRs - Memoize `constraints.dependent_property` as we do with `lazy_property`, since we now consult those properties much more often. ## Tested - [x] added a test for `Dist.support` vs `Dist(**params).support` to ensure static and dynamic attributes agree. - [x] refactoring is covered by existing tests - [x] add test cases for `ReshapedTransform` - [x] add a test for `TransformedDistribution` on a wide grid of input shapes - [x] added a regression test for #34859 cc fehiepsi feynmanliang stefanwebb Pull Request resolved: #50581 Reviewed By: ezyang, glaringlee, jpchen Differential Revision: D26024247 Pulled By: neerajprad fbshipit-source-id: f0b9a296f780ff49659b132409e11a29985dde9b

fritzo added the module: distributions Related to torch.distributions label Jan 15, 2021

fritzo requested a review from neerajprad January 15, 2021 14:53

fritzo commented Jan 15, 2021

View reviewed changes

torch/distributions/transforms.py Show resolved Hide resolved

torch/distributions/transforms.py Outdated Show resolved Hide resolved

fritzo mentioned this pull request Jan 15, 2021

Add a Distribution.infer_shapes() method to statically infer shapes pyro-ppl/pyro#2739

Merged

3 tasks

facebook-github-bot added the cla signed label Jan 15, 2021

pytorchbot added the open source label Jan 15, 2021

fehiepsi reviewed Jan 15, 2021

View reviewed changes

torch/distributions/transformed_distribution.py Outdated Show resolved Hide resolved

fehiepsi reviewed Jan 15, 2021

View reviewed changes

torch/distributions/transforms.py Outdated Show resolved Hide resolved

fritzo mentioned this pull request Jan 17, 2021

Add factory for dynamic TransformOp pyro-ppl/funsor#427

Merged

4 tasks

fritzo commented Jan 17, 2021

View reviewed changes

torch/distributions/transforms.py Show resolved Hide resolved

ngimel added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jan 17, 2021

fehiepsi reviewed Jan 22, 2021

View reviewed changes

torch/distributions/transforms.py Show resolved Hide resolved

fehiepsi mentioned this pull request Jan 22, 2021

Add transforms' forward_shape and inverse_shape pyro-ppl/numpyro#887

Merged

Rebase and squash

4e1a93f

fritzo force-pushed the independent-transform branch from 3ccc720 to 4e1a93f Compare January 22, 2021 15:39

fritzo added 4 commits January 22, 2021 10:57

Fix docstring

65e65d1

Correct formula for D as function of N

70b314b

Fix TransformedDistribution shape logic

a5ed853

Add regression test for #34859

397cbd3

fritzo changed the title ~~Independent transform~~ Fix distribution.Transform shaping logic Jan 22, 2021

fritzo changed the title ~~Fix distribution.Transform shaping logic~~ Fix TransformedDistribution shaping logic Jan 22, 2021

Attempt to pacify mypy

57cbe92

neerajprad reviewed Jan 22, 2021

View reviewed changes

torch/distributions/transforms.py Show resolved Hide resolved

fritzo added 2 commits January 22, 2021 14:37

Address review comment

e55d163

lint

c3f6710

facebook-github-bot reviewed Jan 22, 2021

View reviewed changes

fritzo requested a review from fehiepsi January 22, 2021 19:55

Pacify mypy

e8b121e

fehiepsi reviewed Jan 23, 2021

View reviewed changes

fehiepsi previously approved these changes Jan 23, 2021

View reviewed changes

fehiepsi reviewed Jan 23, 2021

View reviewed changes

neerajprad approved these changes Jan 23, 2021

View reviewed changes

Add SoftmaxTransform.forward_shape(), .inverse_shape()

585caf5

Add test for IndpendentTransform

55a1a92

Add PowerTransform.forward_shape()

f3eebc5

fehiepsi approved these changes Jan 23, 2021

View reviewed changes

facebook-github-bot reviewed Jan 25, 2021

View reviewed changes

facebook-github-bot closed this in facebookresearch/beanmachine@f2fbfaa Jan 26, 2021

fritzo mentioned this pull request Jan 28, 2021

Bump to version 1.5.2 pyro-ppl/pyro#2755

Merged

stefanwebb mentioned this pull request Jan 30, 2021

Adds failing test case for batched base distribution stefanwebb/flowtorch-old#18

Open

fritzo mentioned this pull request Jan 31, 2021

Add Distribution.infer_shapes() for static shape analysis pyro-ppl/numpyro#901

Merged

1 task

fritzo mentioned this pull request Apr 6, 2021

Update pyro.param to read event_dim from the constraint pyro-ppl/pyro#2795

Closed

fritzo mentioned this pull request Jun 14, 2021

Implement a generic Distribution.expand() and use it in a few places #59944

Closed

5 tasks

github-actions bot deleted the independent-transform branch February 10, 2024 01:54

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix TransformedDistribution shaping logic #50581

Fix TransformedDistribution shaping logic #50581

fritzo commented Jan 15, 2021 •

edited

neerajprad commented Jan 22, 2021

neerajprad commented Jan 22, 2021

facebook-github-bot left a comment

codecov bot commented Jan 23, 2021

fehiepsi Jan 23, 2021

fehiepsi left a comment

neerajprad left a comment

fritzo commented Jan 23, 2021 •

edited

fehiepsi commented Jan 23, 2021

facebook-github-bot left a comment

Fix TransformedDistribution shaping logic #50581

Fix TransformedDistribution shaping logic #50581

Conversation

fritzo commented Jan 15, 2021 • edited

Summary of changes

Changes planned for follow-up PRs

Tested

neerajprad commented Jan 22, 2021

neerajprad commented Jan 22, 2021

facebook-github-bot left a comment

Choose a reason for hiding this comment

codecov bot commented Jan 23, 2021

Codecov Report

fehiepsi Jan 23, 2021

Choose a reason for hiding this comment

fehiepsi left a comment

Choose a reason for hiding this comment

neerajprad left a comment

Choose a reason for hiding this comment

fritzo commented Jan 23, 2021 • edited

fehiepsi commented Jan 23, 2021

facebook-github-bot left a comment

Choose a reason for hiding this comment

fritzo commented Jan 15, 2021 •

edited

fritzo commented Jan 23, 2021 •

edited