TransformedDistribution and event_shape

[justindomke]

🐛 Bug

TransformedDistribution does not seem to have correct event_shape when the transform has changed the number of dimensions (e.g. with StickBreakingTransform).

To Reproduce

Steps to reproduce the behavior:

import torch as tt
dist = tt.distributions.Dirichlet(tt.ones(3))
support = dist.support
tform = tt.distributions.constraint_registry.biject_to(support)
dist_unconstrained = tt.distributions.TransformedDistribution(dist,tform.inv)

print(tform)
print(dist_unconstrained.sample())
print(dist_unconstrained.event_shape)

Produces the results of

StickBreakingTransform()
tensor([ 0.3628, -1.9573])
torch.Size([3])

Expected behavior

The last print statement should (seemingly) give:

torch.Size([2])

Environment

• PyTorch Version (e.g., 1.0): 1.1.0
• OS (e.g., Linux): Mac
• How you installed PyTorch (conda, pip, source): conda
• Build command you used (if compiling from source): n/a
• Python version: 3.6.4
• CUDA/cuDNN version: n/a
• GPU models and configuration: n/a
• Any other relevant information: n/a

[alicanb]

@fritzo any idea how to fix this? One idea I have is add a event_shape_modifier to all transforms and replace

pytorch/torch/distributions/transformed_distribution.py

Line 53 in 0745591

shape = self.base_dist.batch_shape + self.base_dist.event_shape

with shape = self.base_dist.batch_shape + torch.Size([s + event_shape_modifier for s in self.base_dist.event_shape])

[fritzo]

It's tough to say without more examples. My first reaction would be to define a method

class Transform(object):
    def transform_event_shape(self, event_shape):
        return event_shape

class StickBreakingTransform(Transform):
    def transform_event_shape(self, event_shape):
        return torch.Size((event_shape[0] - 1,))

and use that in TransformedDistribution. This might not work with jitting.