Adopt strict batch shape semantics for distributions by fritzo · Pull Request #806 · pyro-ppl/pyro

fritzo · 2018-02-23T19:19:30Z

This PR adds checks for stricter use of batch dimensions in Pyro. This is only recently possible since it relies on PyTorch support for scalars.

Why?

Pyro 0.2 will be able to automatically introduce batch dimensions for parallelizing enum_discrete and num_particles. To give Pyro space to create new batch dimensions, we need to restrict how dimensions are introduced by user code.

How?

The new requirements are:

users must declare non-batch dimensions using .reshape(extra_event_dims=n)
batch dimensions will only be allowed inside iarange (and maybe irange?)
to use auto-parallelization, users must request dims via SVI(..., max_iarange_nesting=n)

These requirements are checked by the new helper check_site_shape which uses the new .size field in cond_indep_stack frames.

Tasks

fritzo · 2018-02-27T05:44:40Z

@neerajprad I've merged dev with your #824 into this PR. Lots of tests still fail. Feel free to push commits to this branch. I'll chat you if I start working on it again (at earliest Wednesday).

neerajprad · 2018-02-28T20:11:12Z

+                '- .permute() data dimensions']))
+
+    # Check parallel dimensions on the left of max_iarange_nesting.
+    # TODO


@fritzo : Do you want to add some more checks in this PR?

No, they should be added in later PRs.

neerajprad · 2018-02-28T20:14:54Z

            sigma = torch.pow(self.tau, -0.5)
-            pyro.observe("obs0", dist.LogNormal(mu_latent, sigma), obs=self.data[0])
-            pyro.observe("obs1", dist.LogNormal(mu_latent, sigma), obs=self.data[1])
+            pyro.observe("obs0", dist.LogNormal(mu_latent, sigma), obs=self.data[0].squeeze())


This test fails to meet the threshold if we use a vectorized "obs" inside an iarange.

@martinjankowiak, @fritzo - Is that expected? Is there a reason why the two obs are observed in separate statements?

Just checked with @martinjankowiak that this works as expected inside iarange, so I am not sure why I was seeing a difference. Will make the change and update.

@martinjankowiak can we just delete this flaky expensive test?

yes we can. although maybe keep the transformed_distribution bit?

neerajprad · 2018-02-28T20:17:19Z

            raise NotImplementedError('alpha < 1 is not supported')
        self.alpha = alpha
-        self._standard_gamma = Gamma(alpha, alpha.new([1]).expand_as(alpha))
+        self._standard_gamma = Gamma(alpha, torch.empty_like(alpha).fill_(1).expand_as(alpha))


@fritzo - Had to make a small change to support scalars. Is there a more concise way to do this?

How about alpha.new([1]).squeeze().expand_as(alpha)?

neerajprad · 2018-02-28T20:21:47Z

The times for integration tests seem to have substantially increased with either the checks or the iarange/reshaping operations.

fritzo · 2018-03-01T16:15:28Z

@neeraj, is this ready to merge?

fritzo

LGTM (but we'll need another reviewer since I authored much of this PR)

fritzo · 2018-03-01T16:18:32Z

 @pytest.mark.init(rng_seed=161)
 @pytest.mark.parametrize("batch_size", [3, 5, 7, 8, None])
-@pytest.mark.parametrize("map_type", ["tensor", "list"])
+@pytest.mark.parametrize("map_type", ["iarange", "irange", "range"])


@neerajprad We're actually running more tests in integration_batch_1, so that might be part of the slowdown.

neerajprad · 2018-03-01T17:17:50Z

@martinjankowiak - could you take a look at this PR? It will be nice to get this merged soon before we accumulate merge conflicts.

martinjankowiak · 2018-03-01T17:43:52Z

yeah i'll take a look now

martinjankowiak

lgtm!

my questions are more for my own edification than anything lese

martinjankowiak · 2018-03-01T17:52:17Z

                z_dist = TransformedDistribution(dist.Normal(z_mu, z_sigma), self.iafs)
            else:
                z_dist = dist.Normal(z_mu, z_sigma)
+            assert z_dist.event_shape == ()


when we clean-up the tutorials we should probably put asserts like this everywhere as a way of helping users understand the code

martinjankowiak · 2018-03-01T17:56:16Z

        prior_mu = Variable(torch.zeros([batch_size, self.z_dim]))
        prior_sigma = Variable(torch.ones([batch_size, self.z_dim]))
-        zs = pyro.sample("z", dist.Normal(prior_mu, prior_sigma))
+        zs = pyro.sample("z", dist.Normal(prior_mu, prior_sigma).reshape(extra_event_dims=2))


why extra_event_dims=2? this isn't really necessary is it?

I think that's an oversight on my part. We don't need any reshaping here. Will remove.

martinjankowiak · 2018-03-01T17:57:55Z

        variance = self.get_param("variance").expand_as(f)
-
-        return pyro.sample("y", dist.Normal(f, variance), obs=obs)
+        event_dims = f.dim()


won't event_dims always be 1 here?

I made the change because of the failing gp tutorial, but I'm not too familiar with this code. If f.dim() = 1, we can hard-code that.

fritzo · 2018-03-01T18:52:15Z

        prior_mu = Variable(torch.zeros([batch_size, self.z_dim]))
        prior_sigma = Variable(torch.ones([batch_size, self.z_dim]))
-        zs = pyro.sample("z", dist.Normal(prior_mu, prior_sigma).reshape(extra_event_dims=2))
+        zs = pyro.sample("z", dist.Normal(prior_mu, prior_sigma))


Why don't you need to .reshape(extra_event_dims=1) here?

From my understanding, model_sample is not used for inference but for generating samples separately using the trained decoder, for plotting, etc.

fritzo

LGTM (again, someone else must merge)

fritzo and others added 22 commits February 14, 2018 18:17

Sketch EnumeratePoutine

369ecb1

Merge branch 'dev' into enumerate-parallel

403708b

Fix dimension logic in EnumerateMessenger

95957e6

Add more test examples

ca6da58

Refactor ELBO

c3c0c80

Merge branch 'dev' into enumerate-parallel

db45ca5

Attempt to get batch shapes correct for enum_discrete in trace_elbo

c56cd6a

Merge branch 'dev' into enumerate-parallel

a1f2b14

Simplify Trace_ELBO

c39d44a

Drop special-case for enum_discrete in Trace_ELBO

1b60102

Merge branch 'dev' into enumerate-parallel

ebbe9d1

Replace enum_discrete kwarg with enumerate_discrete() function

49606f8

Completely elimitate enum_discrete kwarg

2a7e25d

Fix bugs in tests/infer/test_enum.py

0da83b9

Rename enumerate_discrete to config_enumerate

45140f4

Merge branch 'dev' into enumerate-parallel

53a44ee

Add analytic KL tests for parallel enumeration

37b6ac9

Add test for sum_rightmost()

5ff4c32

Skip slow tests on travis

7787050

Add another gradient test for enumeration

fe8820c

Add TODOs for more tests

2a52b80

Add failing checks for strict shape semantics

62e2faf

fritzo added the WIP label Feb 23, 2018

fritzo added 7 commits February 23, 2018 13:47

Add variously-sized categoricals test

c61a1ba

Remove excruciatingly slow test

5e9e4ff

Fix scalar error

1947ba0

Flake8

0d695db

Merge branch 'enumerate-parallel' into strict-shape

4015163

Fix zero_grads()

398b044

Get test_valid_models.py to pass tests

a06041f

Merge branch 'dev' into strict-shape

693cf30

This was referenced Feb 26, 2018

working towards nested iarange in tracegraph_elbo #780

Merged

Implement TorchDistribution.mask() for dependent masks #821

Merged

Remove networkx from Trace #822

Closed

fritzo added 2 commits February 26, 2018 18:36

Merge branch 'dev' into strict-shape

f64b51b

Merge branch 'dev' into strict-shape

cd57e9d

fritzo and others added 4 commits February 27, 2018 14:56

Fix shaping errors in dmm.py

7efe00d

fixes to examples and tutorials

b9b631d

fix integration tests

50964cc

fix xfail marker; vae test

c99b9ef

neerajprad reviewed Feb 28, 2018

View reviewed changes

address comment; only use iarange for obs in test_inference

860d8d7

neerajprad added awaiting review and removed WIP labels Mar 1, 2018

fritzo commented Mar 1, 2018

View reviewed changes

martinjankowiak previously approved these changes Mar 1, 2018

View reviewed changes

address comment

180ccb2

neerajprad dismissed martinjankowiak’s stale review via 180ccb2 March 1, 2018 18:26

fritzo commented Mar 1, 2018

View reviewed changes

martinjankowiak approved these changes Mar 1, 2018

View reviewed changes

martinjankowiak merged commit fe32be3 into dev Mar 1, 2018

martinjankowiak deleted the strict-shape branch March 1, 2018 20:17

neerajprad mentioned this pull request Mar 2, 2018

Slow test clean up time #840

Closed

16 tasks

Uh oh!

Conversation

fritzo commented Feb 23, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why?

How?

Tasks

Uh oh!

fritzo commented Feb 27, 2018

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

neerajprad Feb 28, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

neerajprad commented Feb 28, 2018

Uh oh!

fritzo commented Mar 1, 2018

Uh oh!

fritzo left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

neerajprad commented Mar 1, 2018

Uh oh!

martinjankowiak commented Mar 1, 2018

Uh oh!

martinjankowiak left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fritzo left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

fritzo commented Feb 23, 2018 •

edited

Loading

neerajprad Feb 28, 2018 •

edited

Loading