Implement pyro.ops.streaming module

[fritzo]

Addresses #2843

This implements a new module pyro.ops.streaming to streamingly track various statistics. The first intended use case is the planned StreamingMCMC class which will track statistics rather than store samples. There are other potential uses in high-dimensional inference, e.g. recording statistics of gradients during SVI and computing sample moments from predictive when the samples don't fit in memory.

Design choices

The two basic operations are .update() and .get(). The third operation .merge() will be useful for multiple-chain MCMC and computing things like rhat.

I have restricted to the data type to dictionaries of tensors, which is the basic datatype in pyro.infer.mcmc and in much of NumPyro. We could easily generalize this to pytrees by adding classes StatsOfList and StatsOfTuple.

Tested

• tested commutativity of update-get
• tested commutativity and associativity of update-merge-get
• ran mypy locally

[fritzo]

@eb8680 I've added you as a reviewer because these streaming classes create a new semigroup abstraction and you're the resident algebra expert.

[fritzo]

@mtsokol I believe you can use something like the following statistics in #2843:

from pyro.ops.streaming import CountMeanVariance, StatsOfDict

...
stats = StatsOfDict(default=CountMeanVariance)
for mcmc_sample in ...:  # learning loop
    stats.update({
        name: transformed_sample for name, transformed_sample in mcmc_sample.items()
    })
result = stats.get()

Let me know if it looks like you'll need any changes to this PR.

[eb8680]

Neat API! I'm a little confused by some of the types. Did you try running mypy locally?

[eb8680]

nit: these assertions should no longer be necessary with type hints

[fritzo]

Hmm, good point more generally. However I'd like to argue that we should include both type hints and assertions until all common tools can leverage type hints. My reasoning is that I'd really like to catch errors as early as possible, e.g. when users (like me) are working in a jupyter notebook. I think until Jupyter dynamically checks types while editing we'll want extra guard rails especially for tricky interfaces like this.

[eb8680]

Can this be strengthened to Dict[str, Type[StreamingStats]]?

[fritzo]

I strengthened to Hashable, but I think we do want to support e.g. integer keys among chains.

[fritzo]

@eb8680 thanks for reviewing!

I'm a little confused by some of the types. Did you try running mypy locally?

Sorry, I didn't run mypy locally, and some of the types are stale after refactoring. I'll fix... UPDATE ...fixed and ran mypy locally.

[mtsokol]

Current #2857 draft isn't chain-aware and I'm wondering about it. It can be either handled by pyro.ops.streaming, e.g.

class CountMeanStats(StreamingStats):
    def __init__(self, num_chains=1):
        self.counts = [0] * num_chains
        ...

    def update(self, sample, chain_index=0):
        ...

    def get(self, group_by_chain=True):
        # we can sum across chains

so the update in StreamingMCMC would be easy:

self._statistics.update({
    name: transformed_sample for name, transformed_sample in z_acc.items()
}, chain_index)

Otherwise it can be handled by StreamingMCMC with e.g. separate CountMeanStats for each chain that can be returned or merged into one if group_by_chain=False somewhere in summary.
WDYT?

[fritzo]

@mtsokol I think it's best to keep chain logic in the StreamingMCMC class so as to keep StreamingStats subclasses as simple as possible (and hence easy to extend by creating new subclasses). However I think with the latest couple commits you can easily separate by chain by making either a nested StatsOfDict or using keys of the form (chain_id, site["name"]). Let me know if you have any ideas about changing the StreamingStats interface to make this easier.

[eb8680]

LGTM after merge conflicts are resolved

[fritzo]

Thanks for reviewing @eb8680! Looks like I'll be using this right away in my mutation models 😄