Skip to content

Commit

Permalink
[pt2] Add meta for poisson (pytorch#134103)
Browse files Browse the repository at this point in the history
Because aten.poisson doesn't have meta function registered, there is one additional eager execution of this op during compilation phase of torch.compile.

There are more ops without meta registration. Is there any reason for it?
Pull Request resolved: pytorch#134103
Approved by: https://github.com/ezyang
  • Loading branch information
australopitek authored and pytorchmergebot committed Aug 26, 2024
1 parent ed86ac2 commit 38f97ec
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions torch/_meta_registrations.py
Original file line number Diff line number Diff line change
Expand Up @@ -2060,6 +2060,12 @@ def meta_bernoulli_p(self, p=0.5, generator=None):
return torch.empty_like(self).contiguous()


@register_meta([aten.poisson.default, aten.poisson.out])
@out_wrapper()
def meta_poisson(self, generator=None):
return torch.empty_like(self)


@register_meta(aten._fused_moving_avg_obs_fq_helper.default)
def meta__fused_moving_avg_obs_fq_helper(
self,
Expand Down

0 comments on commit 38f97ec

Please sign in to comment.