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.

Sign up for GitHub

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

Only do meta registrations if we have the ops #7500

Merged
merged 4 commits into from
Apr 7, 2023
Merged

Only do meta registrations if we have the ops #7500

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f78e0ad
Only do meta registrations if we have the ops
ezyang Apr 5, 2023
2e12552
fixup
ezyang Apr 6, 2023
f71e298
add import
ezyang Apr 6, 2023
7a58ba1
Merge branch 'main' into optional-reg
ezyang Apr 6, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 10 additions & 6 deletions torchvision/_meta_registrations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import functools

import torch
import torch.library

Expand All @@ -6,20 +8,22 @@

from torch._prims_common import check

_meta_lib = torch.library.Library("torchvision", "IMPL", "Meta")

vision = torch.ops.torchvision
@functools.lru_cache(None)
def get_meta_lib():
return torch.library.Library("torchvision", "IMPL", "Meta")


def register_meta(op):
def register_meta(op_name, overload_name="default"):
def wrapper(fn):
_meta_lib.impl(op, fn)
if torchvision.extension._has_ops():
get_meta_lib().impl(getattr(getattr(torch.ops.torchvision, op_name), overload_name), fn)
return fn

return wrapper


@register_meta(vision.roi_align.default)
@register_meta("roi_align")
def meta_roi_align(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned):
check(rois.size(1) == 5, lambda: "rois must have shape as Tensor[K, 5]")
check(
Expand All @@ -34,7 +38,7 @@ def meta_roi_align(input, rois, spatial_scale, pooled_height, pooled_width, samp
return input.new_empty((num_rois, channels, pooled_height, pooled_width))


@register_meta(vision._roi_align_backward.default)
@register_meta("_roi_align_backward")
def meta_roi_align_backward(
grad, rois, spatial_scale, pooled_height, pooled_width, batch_size, channels, height, width, sampling_ratio, aligned
):
Expand Down