From e782f200f7a525af4d10ff79d066613d46e23356 Mon Sep 17 00:00:00 2001 From: David Riazati Date: Thu, 14 Mar 2019 16:45:31 -0700 Subject: [PATCH] Allow fewer arguments than the max in ArgumentSpec (#17826) Summary: Fixes #17558 The flattened tuple `Optional[Tuple[int, int]]` could either result in 1 (`None`) or 2 (`int` and `int`) values, so allow this case in `ArgumentSpec` Pull Request resolved: https://github.com/pytorch/pytorch/pull/17826 Differential Revision: D14415290 Pulled By: driazati fbshipit-source-id: 971bfa39502cfb8f08a991f16ffed6d138e48dc9 --- test/test_jit.py | 12 ++++++++++++ torch/csrc/jit/argument_spec.h | 2 +- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/test/test_jit.py b/test/test_jit.py index 5725201..baab4d5 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -10636,6 +10636,18 @@ a") m._create_method_from_graph("forward", foo.graph) self.getExportImportCopy(m) + def test_optional_tuple(self): + def fn(x=None): + # type: (Optional[Tuple[int, int]]) -> Tuple[int, int] + if x is None: + new_x = (1, 2) + else: + new_x = x + return new_x + + self.checkScript(fn, ((3, 4),)) + self.checkScript(fn, ()) + def test_split(self): def split_two(tensor): a, b, c = torch.split(tensor, 2, dim=1) diff --git a/torch/csrc/jit/argument_spec.h b/torch/csrc/jit/argument_spec.h index eb3b98d..a345c6b 100644 --- a/torch/csrc/jit/argument_spec.h +++ b/torch/csrc/jit/argument_spec.h @@ -77,7 +77,7 @@ struct ArgumentSpec { for (const auto& i : inputs) { addInput(i, offset, with_grad); } - AT_ASSERT(offset == num_flat_inputs); + AT_ASSERT(offset <= num_flat_inputs); } void addInput(const IValue& input, size_t& offset, bool with_grad) { -- 2.7.4