{"action":"created","comment":{"url":"https://api.github.com/repos/uber/pyro/pulls/comments/186270547","pull_request_review_id":117823111,"id":186270547,"diff_hunk":"@@ -0,0 +1,80 @@\n+import weakref\n+import torch\n+\n+import pyro\n+import pyro.poutine as poutine\n+\n+\n+class compile(object):\n+    \"\"\"\n+    Drop-in replacement for :func:`torch.jit.compile` that works with\n+    Pyro functions that call :func:`pyro.param`.\n+\n+    The actual compilation artifact is stored in ``self.compiled``.\n+    Call diagnostic methods on this attribute.\n+\n+    Example::\n+\n+        def model(x):\n+            scale = pyro.param(\"scale\", torch.tensor(0.5), constraint=constraints.positive)\n+            return pyro.sample(\"y\", dist.Normal(x, scale))\n+\n+        @pyro.ops.jit.compile(nderivs=1)\n+        def model_log_prob_fn(x, y):\n+            cond_model = pyro.condition(model, data={\"y\": y})\n+            tr = pyro.poutine.trace(cond_model).get_trace(x)\n+            return tr.log_prob_sum()\n+    \"\"\"\n+    def __init__(self, **kwargs):\n+        self._kwargs = kwargs\n+        self.compiled = None\n+        self._param_names = None\n+\n+    def __call__(self, fn):\n+\n+        assert self.compiled is None, \\\n+            \"Cannot use a compile instance on more than one function\"\n+\n+        def wrapper(*args, **kwargs):\n+            # if first time\n+            if self.compiled is None:\n+                # param capture\n+                with poutine.block():\n+                    with poutine.trace(param_only=True) as first_param_capture:\n+                        fn(*args, **kwargs)\n+\n+                self._param_names = list(set(first_param_capture.trace.nodes.keys()))\n+\n+                weakself = weakref.ref(self)\n+\n+                @torch.jit.compile(**self._kwargs)\n+                def compiled(unconstrained_params, *args):\n+                    self = weakself()\n+                    constrained_params = {name: pyro.param(name, unconstrained_param)\n+                                          for name, unconstrained_param\n+                                          in zip(self._param_names,\n+                                                 unconstrained_params)}\n+\n+                    return poutine.replay(\n+                        fn, params=constrained_params)(*args, **kwargs)\n+\n+                self.compiled = compiled\n+\n+            param_list = [pyro.param(name).unconstrained()\n+                          for name in self._param_names]\n+\n+            with poutine.block(hide=self._param_names):","path":"pyro/ops/jit.py","position":66,"original_position":66,"commit_id":"dceb1c52bed4cab5980d810adb4c828c3f35a305","original_commit_id":"1d5b750e7d334c8390f3aff50b836fe00b8ea64c","user":{"login":"fritzo","id":648532,"avatar_url":"https://avatars0.githubusercontent.com/u/648532?v=4","gravatar_id":"","url":"https://api.github.com/users/fritzo","html_url":"https://github.com/fritzo","followers_url":"https://api.github.com/users/fritzo/followers","following_url":"https://api.github.com/users/fritzo/following{/other_user}","gists_url":"https://api.github.com/users/fritzo/gists{/gist_id}","starred_url":"https://api.github.com/users/fritzo/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/fritzo/subscriptions","organizations_url":"https://api.github.com/users/fritzo/orgs","repos_url":"https://api.github.com/users/fritzo/repos","events_url":"https://api.github.com/users/fritzo/events{/privacy}","received_events_url":"https://api.github.com/users/fritzo/received_events","type":"User","site_admin":false},"body":"Can you help me undersand why you're blocking and tracing here?","created_at":"2018-05-05T19:51:06Z","updated_at":"2018-05-05T19:52:02Z","html_url":"https://github.com/uber/pyro/pull/1126#discussion_r186270547","pull_request_url":"https://api.github.com/repos/uber/pyro/pulls/1126","author_association":"MEMBER","_links":{"self":{"href":"https://api.github.com/repos/uber/pyro/pulls/comments/186270547"},"html":{"href":"https://github.com/uber/pyro/pull/1126#discussion_r186270547"},"pull_request":{"href":"https://api.github.com/repos/uber/pyro/pulls/1126"}}},"pull_request":{"url":"https://api.github.com/repos/uber/pyro/pulls/1126","id":186172450,"html_url":"https://github.com/uber/pyro/pull/1126","diff_url":"https://github.com/uber/pyro/pull/1126.diff","patch_url":"https://github.com/uber/pyro/pull/1126.patch","issue_url":"https://api.github.com/repos/uber/pyro/issues/1126","number":1126,"state":"open","locked":false,"title":"Add helper decorator for JITing parametrized functions","user":{"login":"eb8680","id":2032320,"avatar_url":"https://avatars0.githubusercontent.com/u/2032320?v=4","gravatar_id":"","url":"https://api.github.com/users/eb8680","html_url":"https://github.com/eb8680","followers_url":"https://api.github.com/users/eb8680/followers","following_url":"https://api.github.com/users/eb8680/following{/other_user}","gists_url":"https://api.github.com/users/eb8680/gists{/gist_id}","starred_url":"https://api.github.com/users/eb8680/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/eb8680/subscriptions","organizations_url":"https://api.github.com/users/eb8680/orgs","repos_url":"https://api.github.com/users/eb8680/repos","events_url":"https://api.github.com/users/eb8680/events{/privacy}","received_events_url":"https://api.github.com/users/eb8680/received_events","type":"User","site_admin":false},"body":"This PR adds a helpful utility `pyro.ops.jit.compile` that serves as a drop-in replacement for `torch.jit.compile` that is compatible with functions that call `pyro.param` internally.  Every such function is implicitly a closure over the Pyro parameter store, but `torch.jit.compile` does not allow closures over tensors.  `pyro.ops.jit.compile` enables this use case and also ensures that parameters are always exposed, even after compilation.\r\n\r\nThis makes the JIT easier to use for common Pyro use cases like computing loss functions in SVI or the energy in HMC that depend on parameters:\r\n```python\r\ndef model(x):\r\n    scale = pyro.param(\"scale\", torch.tensor(0.5), constraint=constraints.positive)\r\n    return pyro.sample(\"y\", dist.Normal(x, scale))\r\n\r\n@pyro.ops.jit.compile(nderivs=1)\r\ndef model_log_prob_fn(x, y):\r\n    cond_model = pyro.condition(model, data={\"y\": y})\r\n    tr = pyro.poutine.trace(cond_model).get_trace(x)\r\n    return tr.log_prob_sum()\r\n\r\n...\r\n\r\n# even after compilation, model_log_prob still exposes params\r\nwith poutine.trace(param_only=True) as param_capture:\r\n    model_log_prob(torch.tensor(1.0), torch.tensor(2.0))\r\n\r\nassert \"p\" in param_capture.trace\r\n```\r\n\r\nThe PR also slightly refactors `JitTrace*_ELBO` to use `pyro.ops.jit.compile` instead of `torch.jit.compile` and modifies `ReplayMessenger` to support replaying of parameters.","created_at":"2018-05-05T19:35:51Z","updated_at":"2018-05-05T19:52:02Z","closed_at":null,"merged_at":null,"merge_commit_sha":"2765315112f7503f5f12c68301f4b7c5ecc7de36","assignee":null,"assignees":[],"requested_reviewers":[],"requested_teams":[],"labels":[{"id":704807577,"url":"https://api.github.com/repos/uber/pyro/labels/WIP","name":"WIP","color":"fbca04","default":false}],"milestone":null,"commits_url":"https://api.github.com/repos/uber/pyro/pulls/1126/commits","review_comments_url":"https://api.github.com/repos/uber/pyro/pulls/1126/comments","review_comment_url":"https://api.github.com/repos/uber/pyro/pulls/comments{/number}","comments_url":"https://api.github.com/repos/uber/pyro/issues/1126/comments","statuses_url":"https://api.github.com/repos/uber/pyro/statuses/dceb1c52bed4cab5980d810adb4c828c3f35a305","head":{"label":"uber:param-jit","ref":"param-jit","sha":"dceb1c52bed4cab5980d810adb4c828c3f35a305","user":{"login":"uber","id":538264,"avatar_url":"https://avatars3.githubusercontent.com/u/538264?v=4","gravatar_id":"","url":"https://api.github.com/users/uber","html_url":"https://github.com/uber","followers_url":"https://api.github.com/users/uber/followers","following_url":"https://api.github.com/users/uber/following{/other_user}","gists_url":"https://api.github.com/users/uber/gists{/gist_id}","starred_url":"https://api.github.com/users/uber/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/uber/subscriptions","organizations_url":"https://api.github.com/users/uber/orgs","repos_url":"https://api.github.com/users/uber/repos","events_url":"https://api.github.com/users/uber/events{/privacy}","received_events_url":"https://api.github.com/users/uber/received_events","type":"Organization","site_admin":false},"repo":{"id":94506832,"name":"pyro","full_name":"uber/pyro","owner":{"login":"uber","id":538264,"avatar_url":"https://avatars3.githubusercontent.com/u/538264?v=4","gravatar_id":"","url":"https://api.github.com/users/uber","html_url":"https://github.com/uber","followers_url":"https://api.github.com/users/uber/followers","following_url":"https://api.github.com/users/uber/following{/other_user}","gists_url":"https://api.github.com/users/uber/gists{/gist_id}","starred_url":"https://api.github.com/users/uber/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/uber/subscriptions","organizations_url":"https://api.github.com/users/uber/orgs","repos_url":"https://api.github.com/users/uber/repos","events_url":"https://api.github.com/users/uber/events{/privacy}","received_events_url":"https://api.github.com/users/uber/received_events","type":"Organization","site_admin":false},"private":false,"html_url":"https://github.com/uber/pyro","description":"Deep universal probabilistic programming with Python and PyTorch","fork":false,"url":"https://api.github.com/repos/uber/pyro","forks_url":"https://api.github.com/repos/uber/pyro/forks","keys_url":"https://api.github.com/repos/uber/pyro/keys{/key_id}","collaborators_url":"https://api.github.com/repos/uber/pyro/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/uber/pyro/teams","hooks_url":"https://api.github.com/repos/uber/pyro/hooks","issue_events_url":"https://api.github.com/repos/uber/pyro/issues/events{/number}","events_url":"https://api.github.com/repos/uber/pyro/events","assignees_url":"https://api.github.com/repos/uber/pyro/assignees{/user}","branches_url":"https://api.github.com/repos/uber/pyro/branches{/branch}","tags_url":"https://api.github.com/repos/uber/pyro/tags","blobs_url":"https://api.github.com/repos/uber/pyro/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/uber/pyro/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/uber/pyro/git/refs{/sha}","trees_url":"https://api.github.com/repos/uber/pyro/git/trees{/sha}","statuses_url":"https://api.github.com/repos/uber/pyro/statuses/{sha}","languages_url":"https://api.github.com/repos/uber/pyro/languages","stargazers_url":"https://api.github.com/repos/uber/pyro/stargazers","contributors_url":"https://api.github.com/repos/uber/pyro/contributors","subscribers_url":"https://api.github.com/repos/uber/pyro/subscribers","subscription_url":"https://api.github.com/repos/uber/pyro/subscription","commits_url":"https://api.github.com/repos/uber/pyro/commits{/sha}","git_commits_url":"https://api.github.com/repos/uber/pyro/git/commits{/sha}","comments_url":"https://api.github.com/repos/uber/pyro/comments{/number}","issue_comment_url":"https://api.github.com/repos/uber/pyro/issues/comments{/number}","contents_url":"https://api.github.com/repos/uber/pyro/contents/{+path}","compare_url":"https://api.github.com/repos/uber/pyro/compare/{base}...{head}","merges_url":"https://api.github.com/repos/uber/pyro/merges","archive_url":"https://api.github.com/repos/uber/pyro/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/uber/pyro/downloads","issues_url":"https://api.github.com/repos/uber/pyro/issues{/number}","pulls_url":"https://api.github.com/repos/uber/pyro/pulls{/number}","milestones_url":"https://api.github.com/repos/uber/pyro/milestones{/number}","notifications_url":"https://api.github.com/repos/uber/pyro/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/uber/pyro/labels{/name}","releases_url":"https://api.github.com/repos/uber/pyro/releases{/id}","deployments_url":"https://api.github.com/repos/uber/pyro/deployments","created_at":"2017-06-16T05:03:47Z","updated_at":"2018-05-05T16:41:23Z","pushed_at":"2018-05-05T19:43:55Z","git_url":"git://github.com/uber/pyro.git","ssh_url":"git@github.com:uber/pyro.git","clone_url":"https://github.com/uber/pyro.git","svn_url":"https://github.com/uber/pyro","homepage":"http://pyro.ai","size":58024,"stargazers_count":3384,"watchers_count":3384,"language":"Jupyter Notebook","has_issues":true,"has_projects":true,"has_downloads":true,"has_wiki":true,"has_pages":true,"forks_count":339,"mirror_url":null,"archived":false,"open_issues_count":45,"license":{"key":"other","name":"Other","spdx_id":null,"url":null},"forks":339,"open_issues":45,"watchers":3384,"default_branch":"dev"}},"base":{"label":"uber:dev","ref":"dev","sha":"ce65866326fca87effbb4171128bb459b167f460","user":{"login":"uber","id":538264,"avatar_url":"https://avatars3.githubusercontent.com/u/538264?v=4","gravatar_id":"","url":"https://api.github.com/users/uber","html_url":"https://github.com/uber","followers_url":"https://api.github.com/users/uber/followers","following_url":"https://api.github.com/users/uber/following{/other_user}","gists_url":"https://api.github.com/users/uber/gists{/gist_id}","starred_url":"https://api.github.com/users/uber/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/uber/subscriptions","organizations_url":"https://api.github.com/users/uber/orgs","repos_url":"https://api.github.com/users/uber/repos","events_url":"https://api.github.com/users/uber/events{/privacy}","received_events_url":"https://api.github.com/users/uber/received_events","type":"Organization","site_admin":false},"repo":{"id":94506832,"name":"pyro","full_name":"uber/pyro","owner":{"login":"uber","id":538264,"avatar_url":"https://avatars3.githubusercontent.com/u/538264?v=4","gravatar_id":"","url":"https://api.github.com/users/uber","html_url":"https://github.com/uber","followers_url":"https://api.github.com/users/uber/followers","following_url":"https://api.github.com/users/uber/following{/other_user}","gists_url":"https://api.github.com/users/uber/gists{/gist_id}","starred_url":"https://api.github.com/users/uber/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/uber/subscriptions","organizations_url":"https://api.github.com/users/uber/orgs","repos_url":"https://api.github.com/users/uber/repos","events_url":"https://api.github.com/users/uber/events{/privacy}","received_events_url":"https://api.github.com/users/uber/received_events","type":"Organization","site_admin":false},"private":false,"html_url":"https://github.com/uber/pyro","description":"Deep universal probabilistic programming with Python and PyTorch","fork":false,"url":"https://api.github.com/repos/uber/pyro","forks_url":"https://api.github.com/repos/uber/pyro/forks","keys_url":"https://api.github.com/repos/uber/pyro/keys{/key_id}","collaborators_url":"https://api.github.com/repos/uber/pyro/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/uber/pyro/teams","hooks_url":"https://api.github.com/repos/uber/pyro/hooks","issue_events_url":"https://api.github.com/repos/uber/pyro/issues/events{/number}","events_url":"https://api.github.com/repos/uber/pyro/events","assignees_url":"https://api.github.com/repos/uber/pyro/assignees{/user}","branches_url":"https://api.github.com/repos/uber/pyro/branches{/branch}","tags_url":"https://api.github.com/repos/uber/pyro/tags","blobs_url":"https://api.github.com/repos/uber/pyro/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/uber/pyro/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/uber/pyro/git/refs{/sha}","trees_url":"https://api.github.com/repos/uber/pyro/git/trees{/sha}","statuses_url":"https://api.github.com/repos/uber/pyro/statuses/{sha}","languages_url":"https://api.github.com/repos/uber/pyro/languages","stargazers_url":"https://api.github.com/repos/uber/pyro/stargazers","contributors_url":"https://api.github.com/repos/uber/pyro/contributors","subscribers_url":"https://api.github.com/repos/uber/pyro/subscribers","subscription_url":"https://api.github.com/repos/uber/pyro/subscription","commits_url":"https://api.github.com/repos/uber/pyro/commits{/sha}","git_commits_url":"https://api.github.com/repos/uber/pyro/git/commits{/sha}","comments_url":"https://api.github.com/repos/uber/pyro/comments{/number}","issue_comment_url":"https://api.github.com/repos/uber/pyro/issues/comments{/number}","contents_url":"https://api.github.com/repos/uber/pyro/contents/{+path}","compare_url":"https://api.github.com/repos/uber/pyro/compare/{base}...{head}","merges_url":"https://api.github.com/repos/uber/pyro/merges","archive_url":"https://api.github.com/repos/uber/pyro/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/uber/pyro/downloads","issues_url":"https://api.github.com/repos/uber/pyro/issues{/number}","pulls_url":"https://api.github.com/repos/uber/pyro/pulls{/number}","milestones_url":"https://api.github.com/repos/uber/pyro/milestones{/number}","notifications_url":"https://api.github.com/repos/uber/pyro/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/uber/pyro/labels{/name}","releases_url":"https://api.github.com/repos/uber/pyro/releases{/id}","deployments_url":"https://api.github.com/repos/uber/pyro/deployments","created_at":"2017-06-16T05:03:47Z","updated_at":"2018-05-05T16:41:23Z","pushed_at":"2018-05-05T19:43:55Z","git_url":"git://github.com/uber/pyro.git","ssh_url":"git@github.com:uber/pyro.git","clone_url":"https://github.com/uber/pyro.git","svn_url":"https://github.com/uber/pyro","homepage":"http://pyro.ai","size":58024,"stargazers_count":3384,"watchers_count":3384,"language":"Jupyter Notebook","has_issues":true,"has_projects":true,"has_downloads":true,"has_wiki":true,"has_pages":true,"forks_count":339,"mirror_url":null,"archived":false,"open_issues_count":45,"license":{"key":"other","name":"Other","spdx_id":null,"url":null},"forks":339,"open_issues":45,"watchers":3384,"default_branch":"dev"}},"_links":{"self":{"href":"https://api.github.com/repos/uber/pyro/pulls/1126"},"html":{"href":"https://github.com/uber/pyro/pull/1126"},"issue":{"href":"https://api.github.com/repos/uber/pyro/issues/1126"},"comments":{"href":"https://api.github.com/repos/uber/pyro/issues/1126/comments"},"review_comments":{"href":"https://api.github.com/repos/uber/pyro/pulls/1126/comments"},"review_comment":{"href":"https://api.github.com/repos/uber/pyro/pulls/comments{/number}"},"commits":{"href":"https://api.github.com/repos/uber/pyro/pulls/1126/commits"},"statuses":{"href":"https://api.github.com/repos/uber/pyro/statuses/dceb1c52bed4cab5980d810adb4c828c3f35a305"}},"author_association":"COLLABORATOR"}}
{"actor":{"display_login":"fritzo"}}