Implicitly broadcast sample sites using iarange dim and size information #1125

neerajprad · 2018-05-04T23:19:31Z

This makes a small change to use the expand logic from #1119 to implement implicit broadcasting via a broadcast poutine.

This should be safe to merge as this logic will not be exercised until the user wraps their model/guides inside poutine.broadcast. We could later decide to reuse this to do the broadcasting behind the scenes for parallelizing over num_particles, for instance. Or use our learnings to build more composable broadcasting effects as @eb8680 mentioned in #1115.

e.g. Using @fritzo's example from #1119

X, Y = 320, 200
x_axis = pyro.iarange("x_axis", X, dim=-2)
y_axis = pyro.iarange("y_axis", Y, dim=-1)
with x_axis:
    x_noise = pyro.sample("x", dist.Normal(0, 1).expand_by([X, 1])) # .expand_by() suffices
with y_axis:
    y_noise = pyro.sample("y", dist.Normal(0, 1).expand_by([Y]))    # .expand_by() suffices
with x_axis, y_axis:
    yx = pyro.sample("yx", dist.Normal(y_noise, 1).expand_by([X])) # .expand_by() suffices
    xy = pyro.sample("xy", dist.Normal(x_noise, 1).expand([X, Y])) # .expand() is needed
    ...

we do not need to expand or expand_by if the model is wrapped inside poutine.broadcast:

x_axis = pyro.iarange("x_axis", X, dim=-2)
y_axis = pyro.iarange("y_axis", Y, dim=-1)
with x_axis:
    x_noise = pyro.sample("x", dist.Normal(0, 1))
    assert x_noise.shape == torch.Size((320, 1))
with y_axis:
    y_noise = pyro.sample("y", dist.Normal(0, 1))
    assert y_noise.shape == torch.Size((200,))
with x_axis, y_axis:
    yx = pyro.sample("yx", dist.Normal(y_noise, 1))
    assert yx.shape == torch.Size((320, 200))
    xy = pyro.sample("xy", dist.Normal(x_noise, 1))
    assert xy.shape == torch.Size((320, 200))

Note that this requires the user to line up the different iaranges correctly via the dim arg.
Tests: Added a couple of tests to test_valid_models. Also modified one of the gradient tests to use the broadcasting logic.

jpchen

is this idempotent? in the sense that if i wrap my models in poutine.broadcast but then dont nest/use broadcasting in my iaranges should everything still work?

jpchen · 2018-05-05T00:14:20Z

pyro/util.py

@@ -218,6 +218,7 @@ def check_site_shape(site, max_iarange_nesting):
    if max_iarange_nesting < len(actual_shape):
        actual_shape = actual_shape[len(actual_shape) - max_iarange_nesting:]

+    expected_shape = broadcast_shape(expected_shape, actual_shape)


we should do strict checking here iiuc that the flag toggles allowing reshaping

neerajprad · 2018-05-05T00:28:37Z

is this idempotent? in the sense that if i wrap my models in poutine.broadcast but then dont nest/use broadcasting in my iaranges should everything still work?

It should work; but this is worth adding as a test.

neerajprad · 2018-05-05T01:27:58Z

Removed generic broadcasting, and added test for idempotence.

fritzo

Nice work, @neerajprad, this is surprisingly simple!

fritzo · 2018-05-06T22:55:23Z

pyro/poutine/broadcast_messenger.py

+    `BroadcastMessenger` automatically broadcasts the batch shape of
+    the stochastic function at a sample site when inside a single
+    or nested iarange context. The existing `batch_shape` must be
+    broadcastable with the size of the :class::`pyro.iarange`


nit:

- :class::`pyro.iarange` + :class:`pyro.iarange`

fritzo · 2018-05-06T23:04:42Z

pyro/poutine/broadcast_messenger.py

+        dist = msg["fn"]
+        actual_batch_shape = getattr(dist, "batch_shape", None)
+        if actual_batch_shape is not None:
+            target_batch_shape = []


Nice! I think this could be a little stricter and a little simpler if you used -1 sizes, something like

target_batch_shape = [-1 if size == 1 else size for size in actual_batch_shape] for f in msg["cond_indep_stack"]: if f.dim is None: continue assert f.dim < 0 if -f.dim > len(target_batch_shape): target_batch_shape = [-1] * (-f.dim - len(target_batch_shape)) + target_batch_shape elif target_batch_shape[f.dim] not in (-1, f.size): raise ValueError("... dim collision ...") target_batch_shape[f.dim] = f.size msg["fn"] = msg["fn"].expand(target_batch_shape)

The problem that I was facing was with expanding smaller-sized latent sites when the iarange dims appear in a staggered fashion, which would cut off the broadcasting starting from the -1 index :

with pyro.iarange("num_particles", 10, dim=-3): with pyro.iarange("components", 2, dim=-1): # with .expand([10, -1, 2]), we get s.shape == (10,) # with .expand([10, 1, 2]), we get s.shape == (10, 1, 2) s = pyro.sample("sample", dist.Bernoulli(0.5)) with pyro.iarange("data", 100, dim=-2): # Note that we need s to have shape (10, 1, 2) here to correctly # expand to (10, 100, 2) ...

Hmm, the snippet I suggested should be order invariant. Note the final .expand() statement is outside of the loop. Am I missing something?

Just updated the snippet to clarify. The problem comes in the sample site in the second iarange where the outermost iarange is at -3 and the inner components one is at -1 (the data iarange is yet to come). If we expand s as dist.Bernoulli(0.5).expand([10, -1, 2]) (note the default -1) it will give us s.shape == torch.Size((10,)) whereas we want it to be of torch.Size((10, 1, 2)) so that it can be correctly broadcasted by the data iarange.

Thanks for explaining, I didn't know that -1 cannot be used for new dimensions in torch.expand(). Could we still try to catch expand errors early, while the iarange.name is still around? Maybe

target_batch_shape = [None if size == 1 else size for size in actual_batch_shape] for f in msg["cond_indep_stack"]: if f.dim is None: continue assert f.dim < 0 target_batch_shape = [None] * (-f.dim - len(target_batch_shape)) + target_batch_shape if target_batch_shape[f.dim] not in (None, f.size): raise ValueError("Shape mismatch inside iarange('{}') at site {} dim {}, {} vs {}".format( f.name, msg['name'], f.dim, f.size, target_batch_shape[dim])) target_batch_shape[f.dim] = f.size # ... remainder of your code ...

Good point; will update.

fritzo · 2018-05-06T23:07:10Z

tests/infer/test_gradient.py

        pyro.sample("nuisance_b", Normal(2, 3))
        pyro.sample("nuisance_a", Normal(0, 1))

    optim = Adam({"lr": 0.1})
+    model, guide = poutine.broadcast(model), poutine.broadcast(guide)


nit: I think it's a little clearer to decorate

@poutine.broadcast def model(): ... @poutine.broadcast def guide(): ...

That way readers know as soon as they start reading the model that it should be read with broadcast semantics.

This is pretty neat!

eb8680

This turned out nice - you should advertise it a little! Maybe add a usage example to the broadcast docstring and a note in the last section of our tensor shape tutorial?

eb8680 · 2018-05-07T07:32:24Z

tests/infer/test_valid_models.py

+    def guide():
+        with pyro.iarange("num_particles", 10, dim=-3):
+            with pyro.iarange("components", 2, dim=-1):
+                pyro.sample("p", dist.Beta(torch.tensor(1.1), torch.tensor(1.1)))


Maybe put this or something similar in the handlers.broadcast docstring as a usage example?

Will update.

neerajprad · 2018-05-07T17:42:18Z

Maybe add a usage example to the broadcast docstring and a note in the last section of our tensor shape tutorial?

Will add to the docstring. I wasn't sure if we should add this to the tutorial yet, because things may change or we may uncover some edge cases, as we discuss and expand our broadcasting semantics to handle different use cases. Let me create a task to update our tutorial once things are stabilized.

eb8680

LGTM

neerajprad · 2018-05-08T15:42:57Z

This should be good to merge, unless there are further comments.

…ion (pyro-ppl#1125)

neerajprad force-pushed the num-particle branch from 59d37a3 to ba50f5b Compare May 4, 2018 23:20

neerajprad requested review from fritzo and eb8680 May 4, 2018 23:20

neerajprad added discussion awaiting review labels May 4, 2018

jpchen reviewed May 5, 2018

View reviewed changes

neerajprad added WIP and removed awaiting review labels May 5, 2018

neerajprad force-pushed the num-particle branch from 9d91232 to 9aa8ebb Compare May 5, 2018 01:03

neerajprad force-pushed the num-particle branch from cae9bfc to ead1576 Compare May 5, 2018 20:18

fritzo reviewed May 6, 2018

View reviewed changes

neerajprad force-pushed the num-particle branch from 7a7cb71 to 5d40cc4 Compare May 7, 2018 06:39

neerajprad added 5 commits May 6, 2018 23:40

Implicitly broadcast sample sites using iarange dim and size information

f03b3b5

address comment

deaa493

clarify comment

f0f747d

lint

2dc7550

address comments

e319851

neerajprad force-pushed the num-particle branch from 5d40cc4 to e319851 Compare May 7, 2018 06:40

eb8680 reviewed May 7, 2018

View reviewed changes

neerajprad added 2 commits May 7, 2018 13:57

address comments; add docs

04949a5

simplify validation

964b10d

neerajprad mentioned this pull request May 8, 2018

Update shapes tutorial with broadcasting utils/changes #1131

Closed

eb8680 approved these changes May 8, 2018

View reviewed changes

neerajprad removed the WIP label May 8, 2018

neerajprad added the awaiting review label May 8, 2018

Merge branch 'dev' into num-particle

b8f77c2

fritzo approved these changes May 9, 2018

View reviewed changes

fritzo merged commit 413e2f2 into pyro-ppl:dev May 9, 2018

neerajprad added a commit to neerajprad/pyro that referenced this pull request May 17, 2018

Implicitly broadcast sample sites using iarange dim and size informat…

7ddccb5

…ion (pyro-ppl#1125)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implicitly broadcast sample sites using iarange dim and size information #1125

Implicitly broadcast sample sites using iarange dim and size information #1125

neerajprad commented May 4, 2018 •

edited

Loading

jpchen left a comment

jpchen May 5, 2018

neerajprad commented May 5, 2018

neerajprad commented May 5, 2018

fritzo left a comment

fritzo May 6, 2018

fritzo May 6, 2018 •

edited

Loading

neerajprad May 7, 2018 •

edited

Loading

fritzo May 7, 2018 •

edited

Loading

neerajprad May 7, 2018

fritzo May 7, 2018 •

edited

Loading

neerajprad May 7, 2018

fritzo May 6, 2018

neerajprad May 7, 2018

eb8680 left a comment

eb8680 May 7, 2018

neerajprad May 7, 2018

neerajprad commented May 7, 2018

eb8680 left a comment

neerajprad commented May 8, 2018

Implicitly broadcast sample sites using iarange dim and size information #1125

Implicitly broadcast sample sites using iarange dim and size information #1125

Conversation

neerajprad commented May 4, 2018 • edited Loading

jpchen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

neerajprad commented May 5, 2018

neerajprad commented May 5, 2018

fritzo left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fritzo May 6, 2018 • edited Loading

Choose a reason for hiding this comment

neerajprad May 7, 2018 • edited Loading

Choose a reason for hiding this comment

fritzo May 7, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fritzo May 7, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eb8680 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

neerajprad commented May 7, 2018

eb8680 left a comment

Choose a reason for hiding this comment

neerajprad commented May 8, 2018

neerajprad commented May 4, 2018 •

edited

Loading

fritzo May 6, 2018 •

edited

Loading

neerajprad May 7, 2018 •

edited

Loading

fritzo May 7, 2018 •

edited

Loading

fritzo May 7, 2018 •

edited

Loading