Skip to content

fix optional type unification #19813

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

Closed
eellison wants to merge 2 commits into from
Closed

fix optional type unification #19813

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5544b14
fix optional unification
Apr 26, 2019
a3a3850
add comment
Apr 26, 2019
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
4 changes: 4 additions & 0 deletions aten/src/ATen/core/jit_type.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,10 @@ using OptionalTypePtr = std::shared_ptr<OptionalType>;
// with a particular type: v: Optional[int] = None()
struct CAFFE2_API OptionalType: public SingleElementType<TypeKind::OptionalType, OptionalType> {
static OptionalTypePtr create(TypePtr element) {
// Optional is a union of [None, T], so Optional[[Optional[T]]] -> Optional[T]
if (auto opt_ptr = element->cast<OptionalType>()) {
return opt_ptr;
}
return OptionalTypePtr(new OptionalType(std::move(element))); // NOLINT(modernize-make-shared)
}
DEFINE_IS_SUBCLASS(OptionalType);
Expand Down
14 changes: 14 additions & 0 deletions test/test_jit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5139,6 +5139,20 @@ def test_while(a, b):
return a + b
''')

def test_opt_opt_refinement(self):
@torch.jit.script
def test_unify(weight, bias):
# type: (Optional[int], Optional[int]) -> Optional[int]
if weight is not None:
opt = None
else:
if bias is not None:
opt = 1
else:
opt = None

return opt

def test_optional_refinement(self):
@torch.jit.script
def test_if_none_assignment(x):
Expand Down