[Core] Fix type checking to be compatible with named tuples (NVIDIA#3986

) * Fix type checking for outputs to be compatible with named tuples Signed-off-by: Vladimir Bataev <[email protected]> * Fix output names in tests Signed-off-by: Vladimir Bataev <[email protected]> Co-authored-by: Somshubra Majumdar <[email protected]> Co-authored-by: Eric Harper <[email protected]>
Zenodia · Apr 14, 2022 · d4408cc · d4408cc
1 parent 2aa6882
commit d4408cc
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 1 deletion.
diff --git a/nemo/core/classes/common.py b/nemo/core/classes/common.py
@@ -262,7 +262,7 @@ def _attach_and_validate_output_types(self, out_objects, ignore_collections=Fals
             mandatory_out_types_list = list(metadata.mandatory_types.items())
 
             # First convert all outputs to list/tuple format to check correct number of outputs
-            if type(out_objects) in (list, tuple):
+            if isinstance(out_objects, (list, tuple)):
                 out_container = out_objects  # can be any rank nested structure
             else:
                 out_container = [out_objects]

diff --git a/tests/core/test_typecheck.py b/tests/core/test_typecheck.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List, NamedTuple
+
 import pytest
 import torch
 
@@ -164,6 +166,32 @@ def __call__(self, x):
         assert result_z.sum() == torch.tensor(20.0)
         assert result_z.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME
 
+    @pytest.mark.unit
+    def test_multiple_output_types_only_namedtuple(self):
+        class NamedTupleOutputType(NamedTuple):
+            y: torch.Tensor
+            z: torch.Tensor
+
+        class MultipleOutputTypesWithNamedTuple(Typing):
+            @property
+            def output_types(self):
+                return {"y": NeuralType(('B',), ElementType()), "z": NeuralType(('B',), ElementType())}
+
+            @typecheck()
+            def __call__(self, x):
+                y = x + 1
+                z = x + 2
+                return NamedTupleOutputType(y=y, z=z)
+
+        obj = MultipleOutputTypesWithNamedTuple()
+        result = obj(x=torch.zeros(10))
+
+        assert result.y.sum() == torch.tensor(10.0)
+        assert result.y.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME
+
+        assert result.z.sum() == torch.tensor(20.0)
+        assert result.z.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME
+
     @pytest.mark.unit
     def test_multiple_mixed_output_types_only(self):
         class MultipleMixedOutputTypes(Typing):
@@ -189,6 +217,35 @@ def __call__(self, x):
         assert result_z[1].sum() == torch.tensor(20.0)
         assert result_z[1].neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME
 
+    @pytest.mark.unit
+    def test_multiple_mixed_output_types_only_namedtuple(self):
+        class NamedTupleOutputType(NamedTuple):
+            y: torch.Tensor
+            zs: List[torch.Tensor]
+
+        class MultipleMixedOutputTypes(Typing):
+            @property
+            def output_types(self):
+                return {"y": NeuralType(('B',), ElementType()), "zs": [NeuralType(('B',), ElementType())]}
+
+            @typecheck()
+            def __call__(self, x):
+                y = x + 1
+                z = x + 2
+                return NamedTupleOutputType(y=y, zs=[z, z])
+
+        obj = MultipleMixedOutputTypes()
+        result_y, result_z = obj(x=torch.zeros(10))
+
+        assert result_y.sum() == torch.tensor(10.0)
+        assert result_y.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME
+
+        assert result_z[0].sum() == torch.tensor(20.0)
+        assert result_z[0].neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME
+
+        assert result_z[1].sum() == torch.tensor(20.0)
+        assert result_z[1].neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME
+
     @pytest.mark.unit
     def test_multiple_mixed_output_types_only_mismatched(self):
         class MultipleMixedOutputTypes(Typing):
@@ -207,6 +264,28 @@ def __call__(self, x):
         with pytest.raises(TypeError):
             result_y, result_z = obj(x=torch.zeros(10))
 
+    @pytest.mark.unit
+    def test_multiple_mixed_output_types_only_namedtuple_mismatched(self):
+        class NamedTupleOutputType(NamedTuple):
+            ys: List[torch.Tensor]
+            z: torch.Tensor
+
+        class MultipleMixedOutputTypes(Typing):
+            @property
+            def output_types(self):
+                return {"ys": NeuralType(('B',), ElementType()), "z": [NeuralType(('B',), ElementType())]}
+
+            @typecheck()
+            def __call__(self, x):
+                # Use list of y, single z, contrary to signature
+                y = x + 1
+                z = x + 2
+                return NamedTupleOutputType(ys=[y, y], z=z)
+
+        obj = MultipleMixedOutputTypes()
+        with pytest.raises(TypeError):
+            _ = obj(x=torch.zeros(10))
+
     @pytest.mark.unit
     def test_incorrect_inheritance(self):
         class IncorrectInheritance(object):