Skip to content

Add a Distribution.infer_shapes() method to statically infer shapes #2739

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 1 commit into from
Feb 23, 2021
Merged

Add a Distribution.infer_shapes() method to statically infer shapes #2739

merged 1 commit into from
Feb 23, 2021

Conversation

fritzo
Copy link
Member

@fritzo fritzo commented Jan 14, 2021 •
edited
Loading

Addresses pyro-ppl/funsor#412
Requires pytorch/pytorch#50547 and pytorch/pytorch#50581

This adds a TorchDistribution.infer_shapes() method to statically compute (batch_shape, event_shape) for use in Funsor. This method inputs *args,**kwargs with the same signature as .__init__() but inputting shapes rather than tensors. The following should be equivalent:

# Eager version.
d = MultivariateNormal(loc, covariance_matrix)
batch_shape = d.batch_shape
event_shape = d.event_shape

# Lazy version.
batch_shape, event_shape = MultivariateNormal.infer_shapes(loc.shape, covariance_matrix.shape)

The default implementation uses new .support.event_dim logic from pytorch/pytorch#50547, so I've included backports in this PR (we can delete next time PyTorch is released).

Tested

  • unit tests for .support.is_discrete
  • unit tests for .support.event_dim
  • unit tests for .infer_shapes()

@fritzo fritzo mentioned this pull request Jan 14, 2021
Closed
@fehiepsi
Copy link
Member

@fritzo Do we need to do the same for transforms?

@fehiepsi fehiepsi mentioned this pull request Jan 14, 2021
Closed
4 tasks
@fritzo
Copy link
Member Author

fritzo commented Jan 15, 2021 •
edited
Loading

Do we need to do the same for transforms?

Indeed to implement .infer_shapes() for TransformedDistribution we will need some sort of static shape computation on Transforms. @stefanwebb has been thinking about this in his Bijectors design doc. I think we will be able to use that doc's suggested .forward_shape() method, but I'd like to wait to implement that until a later PR.

UPDATE I have implemented Transform.forward_shape() and .inverse_shape() in pytorch/pytorch#50581

@fehiepsi fehiepsi self-requested a review January 15, 2021 16:20
This was referenced Jan 17, 2021
Merged
Merged
@fritzo fritzo force-pushed the infer-shapes branch 2 times, most recently from 37d5250 to 3291445 Compare January 27, 2021 20:15
@fritzo fritzo mentioned this pull request Jan 28, 2021
Closed
5 tasks
@fritzo fritzo added the Blocked label Jan 29, 2021
This was referenced Jan 31, 2021
Merged
Closed
@fritzo fritzo mentioned this pull request Feb 11, 2021
Merged
@fritzo fritzo added this to the 1.6 release milestone Feb 11, 2021
@fritzo
Copy link
Member Author

fritzo commented Feb 23, 2021

@fehiepsi this is now unblocked 😄

fehiepsi
fehiepsi approved these changes Feb 23, 2021
View reviewed changes
Copy link
Member

@fehiepsi fehiepsi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@fehiepsi fehiepsi merged commit fd84ea8 into dev Feb 23, 2021
@fritzo
Copy link
Member Author

fritzo commented Feb 23, 2021

Thanks for reviewing @fehiepsi !

@fritzo fritzo mentioned this pull request Feb 23, 2021
Merged
@fritzo fritzo deleted the infer-shapes branch September 27, 2021 14:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@fehiepsi fehiepsi fehiepsi approved these changes

Assignees
No one assigned
Labels
awaiting review enhancement
Projects
None yet
Milestone
1.6 release
Development

Successfully merging this pull request may close these issues.
2 participants
@fritzo @fehiepsi