-
-
Notifications
You must be signed in to change notification settings - Fork 984
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.
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
Added dim
keyword arg
#2472
Added dim
keyword arg
#2472
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM after one nit
We need to take `self.dim` into account at the end of `_call`, `_inverse`, and `log_abs_det_jacobian`. As a clarification, the changes to `log_abs_det_jacobian` are needed to perform an "elementwise" summation, where the elements are the samples/batches. This ensures that the transform works with things like `TransformedDistribution` that use the transform's `event_dim` to sum log probabilities from different transformations.
Use of `event_dim` in `log_abs_det_jacobian` allows `Permute` to be compatible with modules such as `TransformedDistribution` that sum log probabilities.
@akern40 has found a few bugs in how the dimensions are handled, I'll make the changes plus Fritz' requests |
Additional use of `self.dim`
Made the requested changes! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could we add some simple smoke tests to ensure shapes are handled correctly with dim= -1, -2, -3?
@stefanwebb do you want to handle that, or I can? |
@fritzo @akern40 In the process of adding tests I discovered a number of issues with adding a dim keyword to these classes... I had to rethink how tests were handled for I think in an interface redesign I would make the helper functions to create transforms take an All tests pass now! |
This adds a
dim
keyword to thetransforms.Permute
andtransform.AffineCoupling
to define which dimension these transforms operate on. As discussed in #2464 this extra flexibility is useful when dealing with matrix/tensor/etc. valued random variables