Skip to content
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

Tweak automated registration for distribution-to-funsor conversion #391

Merged
merged 4 commits into from
Nov 5, 2020

Conversation

eb8680
Copy link
Member

@eb8680 eb8680 commented Nov 3, 2020

Addresses #386. Motivated by @ordabayevy's forum post about generating Funsor wrappers for custom distributions.

Currently, generic TorchDistribution-to-funsor.distribution.Distribution conversion is handled by a single pattern:

@to_funsor.register(pyro.distributions.Distribution)
def backenddist_to_funsor(backend_dist, output=None, dim_to_name=None):
    funsor_dist_class = ...
    ...

This design requires the pattern to somehow identify a Funsor counterpart funsor_dist_class to backend_dist, which it currently does somewhat hackily by looking in the global namespace. This PR tweaks this logic so that no lookup is necessary:

def backenddist_to_funsor(funsor_dist_class, backend_dist, ...):
    ...

to_funsor.register(MyDistribution)(functools.partial(backenddist_to_funsor, funsor_dist_class))

By including the final line above in the distribution wrapper generator funsor.distribution.make_dist, users should in many cases no longer need to implement to_funsor for custom distributions.

I have also deleted a superfluous conversion pattern funsor.distribution.mvndist_to_funsor whose functionality is duplicated in an eager pattern.

Tested:

  • Refactor exercised by existing tests

@eb8680 eb8680 added enhancement New feature or request refactor labels Nov 3, 2020
@eb8680 eb8680 requested a review from fritzo November 3, 2020 20:19
@eb8680 eb8680 mentioned this pull request Nov 3, 2020
34 tasks
@eb8680 eb8680 mentioned this pull request Nov 4, 2020
3 tasks
@eb8680 eb8680 merged commit 4a2843a into master Nov 5, 2020
@eb8680 eb8680 deleted the automate-conversion-registration branch November 5, 2020 04:38
@eb8680 eb8680 restored the automate-conversion-registration branch November 5, 2020 04:38
@eb8680 eb8680 deleted the automate-conversion-registration branch November 5, 2020 04:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request refactor
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants