chore: import upstream snapshot with attribution

tests/core/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/core/mixins/adapters/test_adapter_mixin.py

@@ -0,0 +1,241 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.core import NeuralModule
+from nemo.core.classes.mixins import adapter_mixin_strategies, adapter_mixins
+from nemo.core.classes.mixins.adapter_mixins import AdapterModuleMixin
+
+
+class DefaultModule(NeuralModule):
+    def __init__(self):
+        super().__init__()
+
+        self.fc = torch.nn.Linear(50, 50)
+        self.bn = torch.nn.BatchNorm1d(50)
+
+    def forward(self, x):
+        x = self.fc(x)
+        x = self.bn(x)
+        out = x
+        return out
+
+    def num_params(self):
+        num: int = 0
+        for p in self.parameters():
+            if p.requires_grad:
+                num += p.numel()
+        return num
+
+
+class DefaultModuleAdapter(DefaultModule, AdapterModuleMixin):
+    def forward(self, x):
+        x = super(DefaultModuleAdapter, self).forward(x)
+
+        if self.is_adapter_available():
+            # For testing purposes, cache the adapter names
+            self._adapter_names = self.get_enabled_adapters()
+            # call forward over model adapters, summing them up
+            x = self.forward_enabled_adapters(x)
+
+        return x
+
+
+def get_adapter_cfg(in_features=50, dim=100, norm_pos='pre'):
+    cfg = {
+        '_target_': 'nemo.collections.common.parts.adapter_modules.LinearAdapter',
+        'in_features': in_features,
+        'dim': dim,
+        'norm_position': norm_pos,
+    }
+    return cfg
+
+
+def get_classpath(cls):
+    return f'{cls.__module__}.{cls.__name__}'
+
+
+if adapter_mixins.get_registered_adapter(DefaultModule) is None:
+    adapter_mixins.register_adapter(DefaultModule, DefaultModuleAdapter)
+
+
+class TestAdapterMixin:
+    @pytest.mark.unit
+    def test_module_registered_adapter_by_class_path(self):
+        classpath = get_classpath(DefaultModule)
+        adapter_meta = adapter_mixins.get_registered_adapter(classpath)
+        assert adapter_meta is not None
+        assert adapter_meta.base_class == DefaultModule
+        assert adapter_meta.adapter_class == DefaultModuleAdapter
+
+    @pytest.mark.unit
+    def test_module_registered_adapter_by_class(self):
+        adapter_meta = adapter_mixins.get_registered_adapter(DefaultModule)
+        assert adapter_meta is not None
+        assert adapter_meta.base_class == DefaultModule
+        assert adapter_meta.adapter_class == DefaultModuleAdapter
+
+    @pytest.mark.unit
+    def test_module_registered_adapter_by_adapter_class(self):
+        adapter_meta = adapter_mixins.get_registered_adapter(DefaultModuleAdapter)
+        assert adapter_meta is not None
+        assert adapter_meta.base_class == DefaultModule
+        assert adapter_meta.adapter_class == DefaultModuleAdapter
+
+    @pytest.mark.unit
+    def test_single_adapter(self):
+        model = DefaultModuleAdapter()
+        original_num_params = model.num_params()
+
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg())
+        new_num_params = model.num_params()
+        assert new_num_params > original_num_params
+
+    @pytest.mark.unit
+    def test_multiple_adapter(self):
+        model = DefaultModuleAdapter()
+        original_num_params = model.num_params()
+
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg())
+        new_num_params = model.num_params()
+        assert new_num_params > original_num_params
+
+        original_num_params = new_num_params
+        model.add_adapter(name='adapter_1', cfg=get_adapter_cfg())
+        new_num_params = model.num_params()
+        assert new_num_params > original_num_params
+
+    @pytest.mark.unit
+    def test_forward_linear_pre(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 50)
+
+        model = DefaultModuleAdapter()
+        origial_output = model(x)
+
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg())
+        new_output = model(x)
+
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.unit
+    def test_forward_linear_post(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 50)
+
+        model = DefaultModuleAdapter()
+        origial_output = model(x)
+
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg(norm_pos='post'))
+        new_output = model(x)
+
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.unit
+    def test_multi_adapter_forward(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 50)
+
+        model = DefaultModuleAdapter()
+        origial_output = model(x)
+
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg())
+        model.add_adapter(name='adapter_1', cfg=get_adapter_cfg())
+        new_output = model(x)
+
+        assert model._adapter_names == ['adapter_0', 'adapter_1']
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.unit
+    def test_multi_adapter_partial_forward(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 50)
+
+        model = DefaultModuleAdapter()
+        origial_output = model(x)
+
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg())
+        model.add_adapter(name='adapter_1', cfg=get_adapter_cfg())
+
+        model.set_enabled_adapters(name='adapter_0', enabled=False)
+        new_output = model(x)
+
+        assert model._adapter_names == ['adapter_1']
+        assert torch.mean(torch.abs(origial_output - new_output)) < 1e-5
+
+    @pytest.mark.unit
+    def test_forward_unfrozen_adapters(self):
+        model = DefaultModuleAdapter()
+        original_num_params = model.num_params()
+
+        dim = 10
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg(dim=dim))
+        model.freeze()
+        model.unfreeze_enabled_adapters()
+
+        assert original_num_params == 2650
+
+        original_params = 0
+        adapter_params = 0
+        for name, param in model.named_parameters():
+            if 'adapter' not in name:
+                assert param.requires_grad is False
+                original_params += param.numel()
+            else:
+                assert param.requires_grad is True
+                adapter_params += param.numel()
+
+        for mname, module in model.named_modules():
+            if isinstance(module, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.BatchNorm3d)):
+                assert module.track_running_stats is False
+
+        assert original_params > adapter_params
+
+    @pytest.mark.unit
+    def test_forward_linear_no_strategy(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 50)
+
+        model = DefaultModuleAdapter()
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg())
+
+        # delete the strategy
+        adapter_module = model.adapter_layer[model.get_enabled_adapters()[0]]
+        del adapter_module.adapter_strategy
+
+        with pytest.raises(AttributeError):
+            _ = model(x)
+
+    @pytest.mark.unit
+    def test_forward_linear_replaced_strategy(self):
+        class MultiplyAdapterStrategy(adapter_mixin_strategies.AbstractAdapterStrategy):
+            def forward(self, input: torch.Tensor, adapter: torch.nn.Module, *, module: AdapterModuleMixin):
+                out = adapter(input)
+                return input * out
+
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 50)
+
+        model = DefaultModuleAdapter()
+        model.add_adapter(name='adapter_0', cfg=get_adapter_cfg())
+
+        # modify the strategy
+        adapter_module = model.adapter_layer[model.get_enabled_adapters()[0]]
+        adapter_module.adapter_strategy = MultiplyAdapterStrategy()
+
+        out = model(x)
+        # result of adapter is zero tensor, output multiplied by adapter result should be zero
+        assert (out > 0.0).any() == torch.tensor(False)

tests/core/mixins/adapters/test_adapter_strategy.py

@@ -0,0 +1,185 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.core import NeuralModule
+from nemo.core.classes.mixins import AdapterModuleMixin, access_mixins, adapter_mixin_strategies, adapter_mixins
+from nemo.utils import config_utils
+
+
+class DefaultModule(NeuralModule):
+    def __init__(self):
+        super().__init__()
+
+        self.fc = torch.nn.Linear(50, 50)
+        self.bn = torch.nn.BatchNorm1d(50)
+
+    def forward(self, x):
+        x = self.fc(x)
+        x = self.bn(x)
+        out = x
+        return out
+
+    def num_params(self):
+        num: int = 0
+        for p in self.parameters():
+            if p.requires_grad:
+                num += p.numel()
+        return num
+
+
+class DefaultModuleAdapter(DefaultModule, AdapterModuleMixin):
+    def forward(self, x):
+        x = super(DefaultModuleAdapter, self).forward(x)
+
+        if self.is_adapter_available():
+            # For testing purposes, cache the adapter names
+            self._adapter_names = self.get_enabled_adapters()
+            # call forward over model adapters, summing them up
+            x = self.forward_enabled_adapters(x)
+
+        return x
+
+
+def get_adapter_cfg(in_features=50, dim=100, norm_pos='pre'):
+    cfg = {
+        '_target_': 'nemo.collections.common.parts.adapter_modules.LinearAdapter',
+        'in_features': in_features,
+        'dim': dim,
+        'norm_position': norm_pos,
+    }
+    return cfg
+
+
+def get_classpath(cls):
+    return f'{cls.__module__}.{cls.__name__}'
+
+
+if adapter_mixins.get_registered_adapter(DefaultModule) is None:
+    adapter_mixins.register_adapter(DefaultModule, DefaultModuleAdapter)
+
+
+class TestAdapterStrategy:
+    @pytest.mark.unit
+    def test_ResidualAddAdapterStrategyConfig(self):
+        IGNORED_ARGS = ['_target_']
+
+        result = config_utils.assert_dataclass_signature_match(
+            adapter_mixin_strategies.ResidualAddAdapterStrategy,
+            adapter_mixin_strategies.ResidualAddAdapterStrategyConfig,
+            ignore_args=IGNORED_ARGS,
+        )
+
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_strategy_default(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 50)
+
+        module = DefaultModuleAdapter()
+        module.add_adapter(name='temp', cfg=get_adapter_cfg())
+        adapter = module.adapter_layer[module.get_enabled_adapters()[0]]
+
+        # update the strategy
+        adapter_strategy = adapter_mixin_strategies.ResidualAddAdapterStrategy()
+        adapter.adapter_strategy = adapter_strategy
+
+        with torch.no_grad():
+            assert adapter_strategy.stochastic_depth == 0.0
+            out = adapter_strategy.forward(x, adapter, module=module)
+            assert (out - x).abs().mean() < 1e-5
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('stochastic_depth', [0.0, 1.0])
+    def test_strategy_stochasic_depth(self, stochastic_depth):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 50)
+
+        module = DefaultModuleAdapter()
+        module.add_adapter(name='temp', cfg=get_adapter_cfg())
+
+        # extract adapter
+        adapter = module.adapter_layer[module.get_enabled_adapters()[0]]
+        # reinitialize the final layer of the adapter module (so that it is not zero init)
+        adapter.module[-1].weight.data += 1
+
+        # get just module output
+        module.set_enabled_adapters('temp', enabled=False)
+        module_out = module(x)
+
+        # get module + adapter output
+        module.set_enabled_adapters('temp', enabled=True)
+        module_adapter_out = module(x)
+
+        assert (
+            module_out - module_adapter_out
+        ).abs().sum() > 0  # results should not be the same after adapter forward now
+
+        adapter_strategy = adapter_mixin_strategies.ResidualAddAdapterStrategy(stochastic_depth=stochastic_depth)
+        adapter.adapter_strategy = adapter_strategy
+
+        module.eval()
+        with torch.inference_mode():  # stochastic depth disabled, no grad tracking
+            assert adapter.adapter_strategy.stochastic_depth == stochastic_depth
+
+            out = adapter_strategy.forward(module_out, adapter, module=module)
+            assert (out - module_adapter_out).abs().mean() < 1e-5
+
+        module.train()
+        with torch.inference_mode():  # stochastic depth enabled, but no grad tracking during training mode
+            out = adapter_strategy.forward(module_out, adapter, module=module)
+
+            if stochastic_depth == 0.0:
+                check = module_adapter_out
+            else:
+                check = module_out
+            assert (out - check).abs().mean() < 1e-5
+
+    @pytest.mark.unit
+    def test_strategy_l2_lambda(self):
+        torch.random.manual_seed(0)
+        x = torch.randn(2, 50)
+
+        module = DefaultModuleAdapter()
+        module.add_adapter(name='temp', cfg=get_adapter_cfg())
+        module.train()
+        adapter = module.adapter_layer[module.get_enabled_adapters()[0]]
+
+        # update the strategy
+        adapter_strategy = adapter_mixin_strategies.ResidualAddAdapterStrategy(l2_lambda=0.01)
+        adapter.adapter_strategy = adapter_strategy
+
+        with torch.no_grad():
+            access_mixins.AccessMixin.reset_registry(module)
+            assert access_mixins.AccessMixin.is_access_enabled() is False
+
+            assert adapter_strategy.stochastic_depth == 0.0
+            assert adapter_strategy.l2_lambda > 0.0
+
+            out = adapter_strategy.forward(x, adapter, module=module)
+            assert (out - x).abs().mean() < 1e-5
+
+            # extract losses
+            assert access_mixins.AccessMixin.is_access_enabled() is True
+            auxiliary_losses = access_mixins.AccessMixin.get_module_registry(module)
+            loss = list(auxiliary_losses.values())[0]
+            assert 'adapter_loss' in loss
+            assert loss['adapter_loss'][0] == torch.tensor(0.0)  # initially adapter is 0 init, no loss required.

tests/core/test_config_utils.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any
+
+import lightning.pytorch as ptl
+import pytest
+from lightning.pytorch.callbacks.early_stopping import EarlyStopping
+
+from nemo.core.config.pytorch_lightning import TrainerConfig
+from nemo.utils import config_utils
+from nemo.utils.exp_manager import EarlyStoppingParams
+
+
+@pytest.fixture()
+def cls():
+    class DummyClass:
+        def __init__(self, a, b=5, c: int = 0, d: 'ABC' = None):
+            pass
+
+    return DummyClass
+
+
+class TestConfigUtils:
+    @pytest.mark.unit
+    def test_all_args_exist(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+            d: Any = None
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_all_args_dont_exist(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert not signatures_match
+        assert len(cls_subset) > 0
+        assert len(dataclass_subset) == 0
+
+    @pytest.mark.unit
+    def test_extra_args_exist(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+            d: Any = None
+            e: float = 0.0
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert not signatures_match
+        assert len(cls_subset) == 0
+        assert len(dataclass_subset) > 0
+
+    @pytest.mark.unit
+    def test_extra_args_exist_but_is_ignored(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+            d: Any = None
+            e: float = 0.0  # Assume ignored
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass, ignore_args=['e'])
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_args_exist_but_is_remapped(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+            e: Any = None  # Assume remapped
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass, remap_args={'e': 'd'})
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_ptl_config(self):
+        PTL_DEPRECATED = []
+        result = config_utils.assert_dataclass_signature_match(ptl.Trainer, TrainerConfig, PTL_DEPRECATED)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_early_stopping_config(
+        self,
+    ):
+        result = config_utils.assert_dataclass_signature_match(EarlyStopping, EarlyStoppingParams)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None

tests/core/test_fault_tolerance.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import lightning.pytorch as pl
+import pytest
+
+from nemo.utils.exp_manager import exp_manager
+
+try:
+    from ptl_resiliency import FaultToleranceCallback
+
+    HAVE_FT = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_FT = False
+
+
+@pytest.mark.skipif(not HAVE_FT, reason="requires resiliency package to be installed.")
+class TestFaultTolerance:
+
+    @pytest.mark.unit
+    def test_fault_tol_callback_not_created_by_default(self):
+        """There should be no FT callback by default"""
+        test_conf = {"create_tensorboard_logger": False, "create_checkpoint_callback": False}
+        test_trainer = pl.Trainer(accelerator='cpu')
+        ft_callback_found = None
+        exp_manager(test_trainer, test_conf)
+        for cb in test_trainer.callbacks:
+            if isinstance(cb, FaultToleranceCallback):
+                ft_callback_found = cb
+        assert ft_callback_found is None
+
+    @pytest.mark.unit
+    def test_fault_tol_callback_created(self):
+        """Verify that fault tolerance callback is created"""
+        try:
+            os.environ['FAULT_TOL_CFG_PATH'] = "/tmp/dummy"
+            test_conf = {
+                "create_tensorboard_logger": False,
+                "create_checkpoint_callback": False,
+                "create_fault_tolerance_callback": True,
+            }
+            test_trainer = pl.Trainer(accelerator='cpu')
+            ft_callback_found = None
+            exp_manager(test_trainer, test_conf)
+            for cb in test_trainer.callbacks:
+                if isinstance(cb, FaultToleranceCallback):
+                    ft_callback_found = cb
+            assert ft_callback_found is not None
+        finally:
+            del os.environ['FAULT_TOL_CFG_PATH']

tests/core/test_fileio.py

@@ -0,0 +1,299 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import tempfile
+
+import numpy as np
+import pytest
+import torch
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.asr.models import EncDecCTCModel
+
+try:
+    from eff.cookbooks import NeMoCookbook
+
+    _EFF_PRESENT_ = True
+except ImportError:
+    _EFF_PRESENT_ = False
+
+# A decorator marking the EFF requirement.
+requires_eff = pytest.mark.skipif(not _EFF_PRESENT_, reason="Export File Format library required to run test")
+
+
+@pytest.fixture()
+def asr_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': 1024,
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
+        'params': {
+            'feat_in': 1024,
+            'num_classes': 28,
+            'vocabulary': [
+                ' ',
+                'a',
+                'b',
+                'c',
+                'd',
+                'e',
+                'f',
+                'g',
+                'h',
+                'i',
+                'j',
+                'k',
+                'l',
+                'm',
+                'n',
+                'o',
+                'p',
+                'q',
+                'r',
+                's',
+                't',
+                'u',
+                'v',
+                'w',
+                'x',
+                'y',
+                'z',
+                "'",
+            ],
+        },
+    }
+    modelConfig = DictConfig(
+        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+    )
+
+    model_instance = EncDecCTCModel(cfg=modelConfig)
+    return model_instance
+
+
+class TestFileIO:
+    @pytest.mark.unit
+    def test_to_from_config_file(self, asr_model):
+        """ " Test makes sure that the second instance created with the same configuration (BUT NOT checkpoint)
+        has different weights."""
+
+        with tempfile.NamedTemporaryFile() as fp:
+            yaml_filename = fp.name
+            asr_model.to_config_file(path2yaml_file=yaml_filename)
+            next_instance = EncDecCTCModel.from_config_file(path2yaml_file=yaml_filename)
+
+            assert isinstance(next_instance, EncDecCTCModel)
+
+            assert len(next_instance.decoder.vocabulary) == 28
+            assert asr_model.num_weights == next_instance.num_weights
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = next_instance.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert not np.array_equal(w1, w2)
+
+    @pytest.mark.unit
+    def test_save_restore_from_nemo_file(self, asr_model):
+        """ " Test makes sure that the second instance created from the same configuration AND checkpoint
+        has the same weights."""
+
+        with tempfile.NamedTemporaryFile() as fp:
+            filename = fp.name
+
+            # Save model (with random artifact).
+            with tempfile.NamedTemporaryFile() as artifact:
+                asr_model.register_artifact(config_path="abc", src=artifact.name)
+                asr_model.save_to(save_path=filename)
+
+            # Restore the model.
+            asr_model2 = EncDecCTCModel.restore_from(restore_path=filename)
+
+            assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary)
+            assert asr_model.num_weights == asr_model2.num_weights
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert np.array_equal(w1, w2)
+
+    @requires_eff
+    @pytest.mark.unit
+    def test_eff_save_restore_from_nemo_file_encrypted(self, asr_model):
+        """ " Test makes sure that after encrypted save-restore the model has the same weights."""
+
+        with tempfile.NamedTemporaryFile() as fp:
+            filename = fp.name
+
+            # Set key - use checkpoint encryption.
+            NeMoCookbook.set_encryption_key("test_key")
+
+            # Save model (with random artifact).
+            with tempfile.NamedTemporaryFile() as artifact:
+                asr_model.register_artifact(config_path="abc", src=artifact.name)
+                asr_model.save_to(save_path=filename)
+
+            # Try to restore the encrypted archive (weights) without the encryption key.
+            NeMoCookbook.set_encryption_key(None)
+            with pytest.raises(PermissionError):
+                # Restore the model.
+                asr_model2 = EncDecCTCModel.restore_from(restore_path=filename)
+
+            # Restore the model.
+            NeMoCookbook.set_encryption_key("test_key")
+            asr_model3 = EncDecCTCModel.restore_from(restore_path=filename)
+            # Reset encryption so it won't mess up with other save/restore.
+            NeMoCookbook.set_encryption_key(None)
+
+            assert asr_model.num_weights == asr_model3.num_weights
+
+    @pytest.mark.unit
+    def test_save_restore_from_nemo_file_with_override(self, asr_model, tmpdir):
+        """ " Test makes sure that the second instance created from the same configuration AND checkpoint
+        has the same weights.
+
+        Args:
+            tmpdir: fixture providing a temporary directory unique to the test invocation.
+        """
+        # Name of the archive in tmp folder.
+        filename = os.path.join(tmpdir, "eff.nemo")
+
+        # Get path where the command is executed - the artifacts will be "retrieved" there.
+        # (original .nemo behavior)
+        cwd = os.getcwd()
+
+        with tempfile.NamedTemporaryFile(mode='a+') as conf_fp:
+
+            # Create a "random artifact".
+            with tempfile.NamedTemporaryFile(mode="w", delete=False) as artifact:
+                artifact.write("magic content 42")
+            # Remember the filename of the artifact.
+            _, artifact_filename = os.path.split(artifact.name)
+            # Add artifact to model.
+            asr_model.register_artifact(config_path="abc", src=artifact.name)
+            # Save model (with "random artifact").
+            asr_model.save_to(save_path=filename)
+
+            # Modify config slightly
+            cfg = asr_model.cfg
+            cfg.encoder.activation = 'swish'
+            yaml_cfg = OmegaConf.to_yaml(cfg)
+            conf_fp.write(yaml_cfg)
+            conf_fp.seek(0)
+
+            # Restore the model.
+            asr_model2 = EncDecCTCModel.restore_from(restore_path=filename, override_config_path=conf_fp.name)
+
+            assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary)
+            assert asr_model.num_weights == asr_model2.num_weights
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert np.array_equal(w1, w2)
+
+            assert asr_model2.cfg.encoder.activation == 'swish'
+
+    @pytest.mark.unit
+    def test_save_model_level_pt_ckpt(self, asr_model):
+        with tempfile.TemporaryDirectory() as ckpt_dir:
+            nemo_file = os.path.join(ckpt_dir, 'asr.nemo')
+            asr_model.save_to(nemo_file)
+
+            # Save model level PT checkpoint
+            asr_model.extract_state_dict_from(nemo_file, ckpt_dir)
+            ckpt_path = os.path.join(ckpt_dir, 'model_weights.ckpt')
+
+            assert os.path.exists(ckpt_path)
+
+            # Restore the model.
+            asr_model2 = EncDecCTCModel.restore_from(restore_path=nemo_file)
+
+            assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary)
+            assert asr_model.num_weights == asr_model2.num_weights
+
+            # Change weights values
+            asr_model2.encoder.encoder[0].mconv[0].conv.weight.data += 1.0
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert not np.array_equal(w1, w2)
+
+            # Restore from checkpoint
+            asr_model2.load_state_dict(torch.load(ckpt_path))
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert np.array_equal(w1, w2)
+
+    @pytest.mark.unit
+    def test_save_module_level_pt_ckpt(self, asr_model):
+        with tempfile.TemporaryDirectory() as ckpt_dir:
+            nemo_file = os.path.join(ckpt_dir, 'asr.nemo')
+            asr_model.save_to(nemo_file)
+
+            # Save model level PT checkpoint
+            asr_model.extract_state_dict_from(nemo_file, ckpt_dir, split_by_module=True)
+            encoder_path = os.path.join(ckpt_dir, 'encoder.ckpt')
+            decoder_path = os.path.join(ckpt_dir, 'decoder.ckpt')
+            preprocessor_path = os.path.join(ckpt_dir, 'preprocessor.ckpt')
+
+            assert os.path.exists(encoder_path)
+            assert os.path.exists(decoder_path)
+            assert os.path.exists(preprocessor_path)
+
+            # Restore the model.
+            asr_model2 = EncDecCTCModel.restore_from(restore_path=nemo_file)
+
+            assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary)
+            assert asr_model.num_weights == asr_model2.num_weights
+
+            # Change weights values
+            asr_model2.encoder.encoder[0].mconv[0].conv.weight.data += 1.0
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert not np.array_equal(w1, w2)
+
+            # Restore from checkpoint
+            asr_model2.encoder.load_state_dict(torch.load(encoder_path))
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert np.array_equal(w1, w2)

tests/core/test_flash_optim.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.core.optim.flash_optim import patch_flashoptim_uneven_shard_support
+
+
+class DummyFlashOptimizer:
+    @staticmethod
+    def _wrap_state_as_dtensor(state, param):
+        state["original"] = True
+
+
+class DummyOptimizer:
+    pass
+
+
+class DummyParam:
+    device_mesh = object()
+    placements = object()
+    shape = torch.Size([5])
+
+    @staticmethod
+    def stride():
+        return (1,)
+
+
+@pytest.mark.unit
+def test_patch_flashoptim_helper_is_noop_for_other_optimizers():
+    optimizer = DummyOptimizer()
+    patch_flashoptim_uneven_shard_support(optimizer)
+    assert not hasattr(DummyOptimizer, "_nemo_patched_uneven_shard")
+
+
+@pytest.mark.unit
+def test_patch_flashoptim_helper_is_idempotent():
+    optimizer = DummyFlashOptimizer()
+    patch_flashoptim_uneven_shard_support(optimizer)
+    patched = DummyFlashOptimizer._wrap_state_as_dtensor
+    patch_flashoptim_uneven_shard_support(DummyFlashOptimizer())
+    assert DummyFlashOptimizer._nemo_patched_uneven_shard is True
+    assert DummyFlashOptimizer._wrap_state_as_dtensor is patched
+
+
+@pytest.mark.unit
+def test_patch_flashoptim_helper_wraps_state_with_shape_and_stride(monkeypatch):
+    calls = {}
+
+    class FakeDTensor:
+        def __init__(self, value=None):
+            self.value = value
+
+        @staticmethod
+        def from_local(local, mesh, placements, shape, stride):
+            calls["local"] = local
+            calls["mesh"] = mesh
+            calls["placements"] = placements
+            calls["shape"] = shape
+            calls["stride"] = stride
+            return FakeDTensor(local)
+
+    import torch.distributed.tensor
+
+    monkeypatch.setattr(torch.distributed.tensor, "DTensor", FakeDTensor)
+
+    optimizer = DummyFlashOptimizer()
+    patch_flashoptim_uneven_shard_support(optimizer)
+
+    state = {"exp_avg": torch.ones(5)}
+    DummyFlashOptimizer._wrap_state_as_dtensor(state, DummyParam())
+
+    assert isinstance(state["exp_avg"], FakeDTensor)
+    assert torch.equal(calls["local"], torch.ones(5))
+    assert calls["shape"] == torch.Size([5])
+    assert calls["stride"] == (1,)

tests/core/test_neural_module.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import tempfile
+
+import pytest
+import torch
+
+from nemo.core.classes.module import NeuralModule
+
+
+class TempModule(NeuralModule):
+
+    def __init__(self):
+        super().__init__()
+
+        self.layer1 = torch.nn.Linear(10, 10, bias=False)
+        self.layer2 = torch.nn.Linear(10, 10, bias=False)
+
+
+class TestNeuralModule:
+
+    @pytest.mark.unit
+    def test_num_weights(self):
+        module = TempModule()
+        assert module.num_weights == 200
+
+    @pytest.mark.unit
+    def test_freeze(self):
+        module = TempModule()
+        module.freeze()
+        for p in module.parameters():
+            assert not p.requires_grad
+
+    @pytest.mark.unit
+    def test_unfreeze(self):
+        module = TempModule()
+        module.freeze()
+        module.unfreeze()
+        for p in module.parameters():
+            assert p.requires_grad
+
+    @pytest.mark.unit
+    def test_as_frozen(self):
+        module = TempModule()
+
+        for p in module.parameters():
+            assert p.requires_grad
+
+        with module.as_frozen():
+            for p in module.parameters():
+                assert not p.requires_grad
+
+        for p in module.parameters():
+            assert p.requires_grad
+
+    @pytest.mark.unit
+    def test_partial_unfreeze(self):
+        module = TempModule()
+
+        for param in module.layer1.parameters():
+            param.requires_grad = False
+
+        module.freeze()
+
+        for param in module.layer1.parameters():
+            assert not param.requires_grad
+
+        assert module._frozen_grad_map is not None
+        assert len(module._frozen_grad_map) == 2
+        assert module._frozen_grad_map['layer1.weight'] is False
+
+        module.unfreeze(partial=True)
+
+        # layer1 should still be frozen due to partial unfreeze
+        assert module.layer1.weight.requires_grad is False
+        assert not hasattr(module, '_frozen_grad_map')

tests/core/test_neural_types.py

@@ -0,0 +1,224 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+
+from nemo.core.neural_types import (
+    AcousticEncodedRepresentation,
+    AudioSignal,
+    AxisKind,
+    AxisKindAbstract,
+    AxisType,
+    ChannelType,
+    ElementType,
+    MelSpectrogramType,
+    MFCCSpectrogramType,
+    NeuralType,
+    NeuralTypeComparisonResult,
+    SpectrogramType,
+    VoidType,
+)
+
+
+class TestNeuralTypeSystem:
+    @pytest.mark.unit
+    def test_short_vs_long_version(self):
+        long_version = NeuralType(
+            axes=(AxisType(AxisKind.Batch, None), AxisType(AxisKind.Dimension, None), AxisType(AxisKind.Time, None)),
+            elements_type=AcousticEncodedRepresentation(),
+        )
+        short_version = NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation())
+        assert long_version.compare(short_version) == NeuralTypeComparisonResult.SAME
+        assert short_version.compare(long_version) == NeuralTypeComparisonResult.SAME
+
+    @pytest.mark.unit
+    def test_parameterized_type_audio_sampling_frequency(self):
+        audio16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000))
+        audio8K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(8000))
+        another16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000))
+
+        assert audio8K.compare(audio16K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+        assert audio16K.compare(audio8K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+        assert another16K.compare(audio16K) == NeuralTypeComparisonResult.SAME
+        assert audio16K.compare(another16K) == NeuralTypeComparisonResult.SAME
+
+    @pytest.mark.unit
+    def test_transpose_same_1(self):
+        type1 = NeuralType(axes=('B', 'T', 'C'))
+        type2 = NeuralType(axes=('T', 'B', 'C'))
+
+        assert type1.compare(type2) == NeuralTypeComparisonResult.TRANSPOSE_SAME
+        assert type2.compare(type1) == NeuralTypeComparisonResult.TRANSPOSE_SAME
+
+    @pytest.mark.unit
+    def test_transpose_same_2(self):
+        audio16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000))
+        audio16K_t = NeuralType(axes=('T', 'B'), elements_type=AudioSignal(16000))
+
+        assert audio16K.compare(audio16K_t) == NeuralTypeComparisonResult.TRANSPOSE_SAME
+
+    @pytest.mark.unit
+    def test_inheritance_spec_augment_example(self):
+        input = NeuralType(('B', 'D', 'T'), SpectrogramType())
+        out1 = NeuralType(('B', 'D', 'T'), MelSpectrogramType())
+        out2 = NeuralType(('B', 'D', 'T'), MFCCSpectrogramType())
+
+        assert out1.compare(out2) == NeuralTypeComparisonResult.INCOMPATIBLE
+        assert out2.compare(out1) == NeuralTypeComparisonResult.INCOMPATIBLE
+        assert input.compare(out1) == NeuralTypeComparisonResult.GREATER
+        assert input.compare(out2) == NeuralTypeComparisonResult.GREATER
+        assert out1.compare(input) == NeuralTypeComparisonResult.LESS
+        assert out2.compare(input) == NeuralTypeComparisonResult.LESS
+
+    @pytest.mark.unit
+    def test_singletone(self):
+        loss_output1 = NeuralType(axes=None)
+        loss_output2 = NeuralType(axes=None)
+
+        assert loss_output1.compare(loss_output2) == NeuralTypeComparisonResult.SAME
+        assert loss_output2.compare(loss_output1) == NeuralTypeComparisonResult.SAME
+
+    @pytest.mark.unit
+    def test_list_of_lists(self):
+        T1 = NeuralType(
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=True),
+                AxisType(kind=AxisKind.Time, size=None, is_list=True),
+                AxisType(kind=AxisKind.Dimension, size=32, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=128, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=256, is_list=False),
+            ),
+            elements_type=ChannelType(),
+        )
+        T2 = NeuralType(
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=False),
+                AxisType(kind=AxisKind.Time, size=None, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=32, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=128, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=256, is_list=False),
+            ),
+            elements_type=ChannelType(),
+        )
+        # TODO: should this be incompatible instead???
+        assert T1.compare(T2), NeuralTypeComparisonResult.TRANSPOSE_SAME
+
+    @pytest.mark.unit
+    def test_void(self):
+        btc_spctr = NeuralType(('B', 'T', 'C'), SpectrogramType())
+        btc_spct_bad = NeuralType(('B', 'T'), SpectrogramType())
+        btc_void = NeuralType(('B', 'T', 'C'), VoidType())
+
+        assert btc_void.compare(btc_spctr) == NeuralTypeComparisonResult.SAME
+        assert btc_spctr.compare(btc_void) == NeuralTypeComparisonResult.INCOMPATIBLE
+        assert btc_void.compare(btc_spct_bad) == NeuralTypeComparisonResult.INCOMPATIBLE
+
+    @pytest.mark.unit
+    def test_big_void(self):
+        big_void_1 = NeuralType(elements_type=VoidType())
+        big_void_2 = NeuralType()
+
+        btc_spctr = NeuralType(('B', 'T', 'C'), SpectrogramType())
+        btc_spct_bad = NeuralType(('B', 'T'), SpectrogramType())
+        t1 = NeuralType(
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=True),
+                AxisType(kind=AxisKind.Time, size=None, is_list=True),
+                AxisType(kind=AxisKind.Dimension, size=32, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=128, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=256, is_list=False),
+            ),
+            elements_type=ChannelType(),
+        )
+        t2 = NeuralType(
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=False),
+                AxisType(kind=AxisKind.Time, size=None, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=32, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=128, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=256, is_list=False),
+            ),
+            elements_type=ChannelType(),
+        )
+
+        assert big_void_1.compare(btc_spctr) == NeuralTypeComparisonResult.SAME
+        assert big_void_1.compare(btc_spct_bad) == NeuralTypeComparisonResult.SAME
+        assert big_void_1.compare(t1) == NeuralTypeComparisonResult.SAME
+        assert big_void_1.compare(t2) == NeuralTypeComparisonResult.SAME
+
+        assert big_void_2.compare(btc_spctr) == NeuralTypeComparisonResult.SAME
+        assert big_void_2.compare(btc_spct_bad) == NeuralTypeComparisonResult.SAME
+        assert big_void_2.compare(t1) == NeuralTypeComparisonResult.SAME
+        assert big_void_2.compare(t2) == NeuralTypeComparisonResult.SAME
+
+    @pytest.mark.unit
+    def test_unspecified_dimensions(self):
+        t0 = NeuralType(
+            (AxisType(AxisKind.Batch, 64), AxisType(AxisKind.Time, 10), AxisType(AxisKind.Dimension, 128)),
+            SpectrogramType(),
+        )
+        t1 = NeuralType(('B', 'T', 'C'), SpectrogramType())
+
+        assert t1.compare(t0), NeuralTypeComparisonResult.SAME
+        assert t0.compare(t1), NeuralTypeComparisonResult.DIM_INCOMPATIBLE
+
+    @pytest.mark.unit
+    def test_any_axis(self):
+        t0 = NeuralType(('B', 'Any', 'Any'), VoidType())
+        t1 = NeuralType(('B', 'Any', 'Any'), SpectrogramType())
+        t2 = NeuralType(('B', 'T', 'C'), SpectrogramType())
+
+        assert t0.compare(t1) == NeuralTypeComparisonResult.SAME
+        assert t0.compare(t2) == NeuralTypeComparisonResult.SAME
+        assert t1.compare(t2) == NeuralTypeComparisonResult.SAME
+        assert t2.compare(t1) == NeuralTypeComparisonResult.INCOMPATIBLE
+        assert t1.compare(t0) == NeuralTypeComparisonResult.INCOMPATIBLE
+
+    @pytest.mark.unit
+    def test_struct(self):
+        class BoundingBox(ElementType):
+            def __str__(self):
+                return "bounding box from detection model"
+
+            def fields(self):
+                return ("X", "Y", "W", "H")
+
+        # ALSO ADD new, user-defined, axis kind
+        class AxisKind2(AxisKindAbstract):
+            Image = 0
+
+        T1 = NeuralType(
+            elements_type=BoundingBox(),
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=True),
+                AxisType(kind=AxisKind2.Image, size=None, is_list=True),
+            ),
+        )
+
+        class BadBoundingBox(ElementType):
+            def __str__(self):
+                return "bad bounding box from detection model"
+
+            def fields(self):
+                return ("X", "Y", "H")
+
+        T2 = NeuralType(
+            elements_type=BadBoundingBox(),
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=True),
+                AxisType(kind=AxisKind2.Image, size=None, is_list=True),
+            ),
+        )
+        assert T2.compare(T1) == NeuralTypeComparisonResult.INCOMPATIBLE

tests/core/test_safe_instantiate.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import patch
+
+import pytest
+import torch
+from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.strategies import DDPStrategy
+from omegaconf import DictConfig
+from torch.distributed.fsdp import MixedPrecisionPolicy
+
+from nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers import BaseTokenizer
+from nemo.collections.tts.g2p.models.base import BaseG2p
+from nemo.core.classes.common import _is_target_allowed, safe_instantiate
+from nemo.utils.decorators import experimental
+
+
+class MockDataset(torch.utils.data.Dataset):
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, index):
+        return index
+
+
+def get_class_path(cls):
+    return f"{cls.__module__}.{cls.__name__}"
+
+
+class MockG2p(BaseG2p):
+    def __call__(self, text: str) -> str:
+        return text
+
+
+class MockTokenizer(BaseTokenizer):
+    def __init__(self):
+        super().__init__(tokens=["a"])
+
+    def encode(self, text: str) -> list[int]:
+        return [0]
+
+
+@experimental
+class MockExperimentalModule(torch.nn.Module):
+    """nn.Module wrapped by @experimental (wrapt), like asr's TransformerEncoder."""
+
+    def __init__(self, value: int = 0):
+        super().__init__()
+        self.value = value
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "config,expected_type",
+    [
+        ({"_target_": "torch.nn.Linear", "in_features": 1, "out_features": 1}, torch.nn.Linear),
+        ({"_target_": get_class_path(MockDataset)}, MockDataset),
+        ({"_target_": "torch.distributed.fsdp.MixedPrecisionPolicy"}, MixedPrecisionPolicy),
+        ({"_target_": "lightning.pytorch.callbacks.ModelCheckpoint"}, ModelCheckpoint),
+        ({"_target_": "lightning.pytorch.strategies.DDPStrategy"}, DDPStrategy),
+    ],
+)
+def test_safe_instantiate_allows_approved_targets(config, expected_type):
+    obj = safe_instantiate(DictConfig(config))
+    assert isinstance(obj, expected_type)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "target,target_type",
+    [
+        ("nemo_text_processing.text_normalization.normalize.Normalizer", type("MockNormalizer", (), {})),
+        ("nemo.collections.tts.torch.g2ps.EnglishG2p", MockG2p),
+        ("nemo.collections.tts.torch.tts_tokenizers.EnglishPhonemesTokenizer", MockTokenizer),
+    ],
+)
+def test_safe_instantiate_allows_exact_exception_targets(target, target_type):
+    sentinel = object()
+    config = DictConfig({"_target_": target})
+
+    with (
+        patch("hydra.utils.get_class", return_value=target_type),
+        patch("hydra.utils.instantiate", return_value=sentinel) as instantiate_mock,
+    ):
+        obj = safe_instantiate(config)
+
+    assert obj is sentinel
+    instantiate_mock.assert_called_once_with(config)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "target",
+    [
+        "subprocess.Popen",
+        "builtins.open",
+        "os.system",
+    ],
+)
+def test_safe_instantiate_blocks_unsafe_targets_before_hydra(target):
+    config = DictConfig({"_target_": target})
+    with patch("hydra.utils.instantiate") as instantiate_mock:
+        with pytest.raises(ValueError, match=f"Instantiation of unsafe target '{target}' is blocked"):
+            safe_instantiate(config)
+
+    instantiate_mock.assert_not_called()
+
+
+@pytest.mark.unit
+def test_safe_instantiate_validates_nested_targets_before_hydra():
+    config = DictConfig(
+        {
+            "_target_": "torch.nn.ModuleList",
+            "modules": [
+                {"_target_": "torch.nn.Linear", "in_features": 1, "out_features": 1},
+                {"_target_": "subprocess.Popen"},
+            ],
+        }
+    )
+
+    with patch("hydra.utils.instantiate") as instantiate_mock:
+        with pytest.raises(ValueError, match="Instantiation of unsafe target 'subprocess.Popen' is blocked"):
+            safe_instantiate(config)
+
+    instantiate_mock.assert_not_called()
+
+
+@pytest.mark.unit
+def test_safe_instantiate_allows_wrapt_decorated_module():
+    """Regression: a wrapt-decorated (@experimental) nn.Module must not be blocked."""
+    target = get_class_path(MockExperimentalModule)
+
+    assert _is_target_allowed(target) is True
+
+    obj = safe_instantiate(DictConfig({"_target_": target, "value": 7}))
+    assert isinstance(obj, torch.nn.Module)
+    assert obj.value == 7
+
+
+@pytest.mark.unit
+def test_safe_instantiate_allows_experimental_asr_transformer_encoder():
+    """The originally reported target: an @experimental nn.Module in asr.modules."""
+    pytest.importorskip("nemo.collections.asr.modules.transformer_encoder")
+
+    assert _is_target_allowed("nemo.collections.asr.modules.transformer_encoder.TransformerEncoder") is True
+    assert _is_target_allowed("nemo.collections.asr.modules.TransformerEncoder") is True

tests/core/test_safe_instantiate_usage.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ast
+import warnings
+from pathlib import Path
+
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+SAFE_INSTANTIATE_WRAPPER = REPO_ROOT / "nemo/core/classes/common.py"
+SOURCE_DIRS = ("nemo", "scripts", "examples")
+
+
+def _iter_python_files():
+    for source_dir in SOURCE_DIRS:
+        yield from (REPO_ROOT / source_dir).rglob("*.py")
+
+
+def _name_for_call(node: ast.AST) -> str | None:
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.Attribute):
+        parent = _name_for_call(node.value)
+        if parent is not None:
+            return f"{parent}.{node.attr}"
+    return None
+
+
+class DirectHydraInstantiateVisitor(ast.NodeVisitor):
+    def __init__(self, path: Path):
+        self.path = path
+        self.function_stack = []
+        self.hydra_aliases = set()
+        self.hydra_utils_aliases = set()
+        self.hydra_instantiate_aliases = set()
+        self.violations = []
+
+    def visit_Import(self, node: ast.Import):
+        for alias in node.names:
+            if alias.name == "hydra":
+                self.hydra_aliases.add(alias.asname or "hydra")
+            elif alias.name == "hydra.utils" and alias.asname:
+                self.hydra_utils_aliases.add(alias.asname)
+        self.generic_visit(node)
+
+    def visit_ImportFrom(self, node: ast.ImportFrom):
+        if node.module == "hydra":
+            for alias in node.names:
+                if alias.name == "utils":
+                    self.hydra_utils_aliases.add(alias.asname or alias.name)
+        elif node.module == "hydra.utils":
+            for alias in node.names:
+                if alias.name == "instantiate":
+                    self.hydra_instantiate_aliases.add(alias.asname or alias.name)
+        self.generic_visit(node)
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        self.function_stack.append(node.name)
+        self.generic_visit(node)
+        self.function_stack.pop()
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
+        self.function_stack.append(node.name)
+        self.generic_visit(node)
+        self.function_stack.pop()
+
+    def visit_Call(self, node: ast.Call):
+        call_name = _name_for_call(node.func)
+        if self._is_direct_hydra_instantiate(call_name) and not self._is_allowed_wrapper_call():
+            self.violations.append((node.lineno, call_name))
+        self.generic_visit(node)
+
+    def _is_direct_hydra_instantiate(self, call_name: str | None) -> bool:
+        if call_name is None:
+            return False
+        if call_name in self.hydra_instantiate_aliases:
+            return True
+        if call_name.endswith(".instantiate"):
+            prefix = call_name[: -len(".instantiate")]
+            if prefix in self.hydra_utils_aliases:
+                return True
+            return any(prefix == f"{hydra_alias}.utils" for hydra_alias in self.hydra_aliases)
+        return False
+
+    def _is_allowed_wrapper_call(self) -> bool:
+        return self.path == SAFE_INSTANTIATE_WRAPPER and self.function_stack[-1:] == ["safe_instantiate"]
+
+
+@pytest.mark.unit
+def test_hydra_instantiate_is_only_called_by_safe_instantiate():
+    violations = []
+    for path in _iter_python_files():
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", SyntaxWarning)
+            tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+        visitor = DirectHydraInstantiateVisitor(path)
+        visitor.visit(tree)
+        violations.extend((path.relative_to(REPO_ROOT), lineno, call_name) for lineno, call_name in visitor.violations)
+
+    assert (
+        not violations
+    ), "Use nemo.core.classes.common.safe_instantiate instead of hydra.utils.instantiate:\n" + "\n".join(
+        f"{path}:{lineno} calls {call_name}" for path, lineno, call_name in violations
+    )

tests/core/test_serialization.py

@@ -0,0 +1,194 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import patch
+
+import pytest
+from omegaconf import DictConfig
+
+from nemo.collections.asr.models import EncDecCTCModel
+from nemo.collections.asr.modules import SpectrogramAugmentation
+from nemo.core.classes.common import Serialization
+
+
+def get_class_path(cls):
+    return f"{cls.__module__}.{cls.__name__}"
+
+
+class MockSerializationImpl(Serialization):
+    def __init__(self, cfg: DictConfig):
+        self.cfg = cfg
+        self.value = self.__class__.__name__
+
+
+class MockSerializationImplV2(MockSerializationImpl):
+    pass
+
+
+class TestSerialization:
+    @pytest.mark.unit
+    def test_from_config_dict_with_cls(self):
+        """Here we test that instantiation works for configs with cls class path in them.
+        Note that just Serialization.from_config_dict can be used to create an object"""
+        config = DictConfig(
+            {
+                'cls': 'nemo.collections.asr.modules.SpectrogramAugmentation',
+                'params': {
+                    'rect_freq': 50,
+                    'rect_masks': 5,
+                    'rect_time': 120,
+                },
+            }
+        )
+        obj = Serialization.from_config_dict(config=config)
+        assert isinstance(obj, SpectrogramAugmentation)
+
+    @pytest.mark.unit
+    def test_from_config_dict_without_cls(self):
+        """Here we test that instantiation works for configs without cls class path in them.
+        IMPORTANT: in this case, correct class type should call from_config_dict. This should work for Models."""
+        preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+        encoder = {
+            'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+            'params': {
+                'feat_in': 64,
+                'activation': 'relu',
+                'conv_mask': True,
+                'jasper': [
+                    {
+                        'filters': 1024,
+                        'repeat': 1,
+                        'kernel': [1],
+                        'stride': [1],
+                        'dilation': [1],
+                        'dropout': 0.0,
+                        'residual': False,
+                        'separable': True,
+                        'se': True,
+                        'se_context_size': -1,
+                    }
+                ],
+            },
+        }
+
+        decoder = {
+            'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
+            'params': {
+                'feat_in': 1024,
+                'num_classes': 28,
+                'vocabulary': [
+                    ' ',
+                    'a',
+                    'b',
+                    'c',
+                    'd',
+                    'e',
+                    'f',
+                    'g',
+                    'h',
+                    'i',
+                    'j',
+                    'k',
+                    'l',
+                    'm',
+                    'n',
+                    'o',
+                    'p',
+                    'q',
+                    'r',
+                    's',
+                    't',
+                    'u',
+                    'v',
+                    'w',
+                    'x',
+                    'y',
+                    'z',
+                    "'",
+                ],
+            },
+        }
+        modelConfig = DictConfig(
+            {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+        )
+        obj = EncDecCTCModel.from_config_dict(config=modelConfig)
+        assert isinstance(obj, EncDecCTCModel)
+
+    @pytest.mark.unit
+    def test_config_updated(self):
+        config = DictConfig(
+            {
+                'cls': 'nemo.collections.asr.modules.SpectrogramAugmentation',
+                'params': {
+                    'rect_freq': 50,
+                    'rect_masks': 5,
+                    'rect_time': 120,
+                },
+            }
+        )
+        obj = Serialization.from_config_dict(config=config)
+        new_config = obj.to_config_dict()
+        assert config != new_config
+        assert 'params' not in new_config
+        assert 'cls' not in new_config
+        assert '_target_' in new_config
+
+    @pytest.mark.unit
+    def test_base_class_instantiation(self):
+        # Target class is V2 impl, calling class is Serialization (base class)
+        config = DictConfig({'target': get_class_path(MockSerializationImplV2)})
+        obj = Serialization.from_config_dict(config=config)
+        new_config = obj.to_config_dict()
+        assert config == new_config
+        assert isinstance(obj, MockSerializationImplV2)
+        assert obj.value == "MockSerializationImplV2"
+
+    @pytest.mark.unit
+    def test_unsafe_legacy_target_is_rejected_before_import(self):
+        config = DictConfig({'target': 'subprocess.Popen'})
+        with patch('nemo.core.classes.common.import_class_by_path') as import_class_mock:
+            with pytest.raises(ValueError, match="Instantiation of unsafe target 'subprocess.Popen' is blocked"):
+                Serialization.from_config_dict(config=config)
+
+        import_class_mock.assert_not_called()
+
+    @pytest.mark.unit
+    def test_legacy_model_support_target_falls_back_to_calling_class(self):
+        config = DictConfig({'target': 'src.multi_classification_models.EncDecMultiClassificationModel'})
+        with patch('nemo.core.classes.common.import_class_by_path') as import_class_mock:
+            obj = MockSerializationImpl.from_config_dict(config=config)
+
+        import_class_mock.assert_not_called()
+        assert isinstance(obj, MockSerializationImpl)
+        assert obj.value == "MockSerializationImpl"
+
+    @pytest.mark.unit
+    def test_self_class_instantiation(self):
+        # Target class is V1 impl, calling class is V1 (same class)
+        config = DictConfig({'target': get_class_path(MockSerializationImpl)})
+        obj = MockSerializationImpl.from_config_dict(config=config)  # Serialization is base class
+        new_config = obj.to_config_dict()
+        assert config == new_config
+        assert isinstance(obj, MockSerializationImpl)
+        assert obj.value == "MockSerializationImpl"
+
+    @pytest.mark.unit
+    def test_sub_class_instantiation(self):
+        # Target class is V1 impl, calling class is V2 (sub class)
+        config = DictConfig({'target': get_class_path(MockSerializationImpl)})
+        obj = MockSerializationImplV2.from_config_dict(config=config)  # Serialization is base class
+        new_config = obj.to_config_dict()
+        assert config == new_config
+        assert isinstance(obj, MockSerializationImplV2)
+        assert obj.value == "MockSerializationImplV2"

tests/core/test_torch_jit_script.py

@@ -0,0 +1,191 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from inspect import isclass
+from typing import Any, Type
+
+import pytest
+import torch
+import torch.nn as nn
+
+import nemo.core.neural_types.elements as nelements
+from nemo.core.classes import Exportable, NeuralModule, typecheck
+from nemo.core.classes.mixins import AdapterModuleMixin
+from nemo.core.neural_types import FloatType, NeuralType, VoidType
+
+
+def get_all_neural_types() -> list[Type[nelements.ElementType]]:
+    """Get all neural types (elements) by inspecting neural_types.element module"""
+    neural_types = []
+    for neural_type_str in dir(nelements):
+        candidate = getattr(nelements, neural_type_str)
+        if (
+            isclass(candidate)
+            and issubclass(candidate, nelements.ElementType)
+            and candidate is not nelements.ElementType
+        ):
+            neural_types.append(candidate)
+    return neural_types
+
+
+class SimpleLinear(NeuralModule):
+    """Simple linear projection. Test use of NeuralModule instead of nn.Module"""
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 20, bias=False)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class SimpleLinearExportable(NeuralModule, Exportable):
+    """Simple linear projection. Test use of NeuralModule with Exportable mixin"""
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 20, bias=False)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class SimpleLinearWithAdapterMixin(NeuralModule, AdapterModuleMixin):
+    """Simple linear projection. Test use of NeuralModule with AdapterModuleMixin mixin"""
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 20, bias=False)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class SimpleLinearWithTypes(NeuralModule):
+    """Simple linear projection. Test use of NeuralModule with input/output types"""
+
+    @property
+    def input_types(self) -> dict[str, Any]:
+        return {
+            "x": NeuralType(None, FloatType()),
+        }
+
+    @property
+    def output_types(self) -> dict[str, Any]:
+        return {
+            "output": NeuralType(None, VoidType()),
+        }
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 20, bias=False)
+
+    @typecheck()
+    def forward(self, x):
+        return self.linear(x)
+
+
+class DummyModuleWithIOTypes(NeuralModule):
+    """Identity module. For testing input/output types in a sequence of network calls"""
+
+    @property
+    def input_types(self) -> dict[str, NeuralType]:
+        return {
+            "x": NeuralType(None, VoidType()),
+        }
+
+    @property
+    def output_types(self) -> dict[str, NeuralType]:
+        return {
+            "output": NeuralType(None, FloatType()),
+        }
+
+    def forward(self, x):
+        return x
+
+
+class TestTorchJitCompatibility:
+    @pytest.mark.unit
+    def test_simple_linear(self):
+        """Test basic module derived from NeuralModule"""
+        module = torch.jit.script(SimpleLinear())
+        x = torch.zeros(2, 10)
+        result = module(x)
+        assert result.shape == (2, 20)
+        assert torch.allclose(result, torch.zeros_like(result))
+
+    @pytest.mark.unit
+    def test_simple_linear_exportable(self):
+        """Test basic module derived from NeuralModule and Exportable"""
+        module = torch.jit.script(SimpleLinearExportable())
+        x = torch.zeros(2, 10)
+        result = module(x)
+        assert result.shape == (2, 20)
+        assert torch.allclose(result, torch.zeros_like(result))
+
+    @pytest.mark.unit
+    def test_simple_linear_with_adapter_mixin(self):
+        """Test basic module derived from NeuralModule and Adapter Mixin"""
+        module = torch.jit.script(SimpleLinearWithAdapterMixin())
+        x = torch.zeros(2, 10)
+        result = module(x)
+        assert result.shape == (2, 20)
+        assert torch.allclose(result, torch.zeros_like(result))
+
+    @pytest.mark.unit
+    def test_simple_linear_with_types(self):
+        """Test basic module derived from NeuralModule containing types"""
+        module = torch.jit.script(SimpleLinearWithTypes())
+        x = torch.zeros(2, 10)
+        result = module(x)
+        assert result.shape == (2, 20)
+        assert torch.allclose(result, torch.zeros_like(result))
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "neural_type",
+        get_all_neural_types(),
+    )
+    def test_element_compilable(self, neural_type: Type[nelements.ElementType]):
+        """
+        Tests that all NeuralType Elements are compilable by TorchScript (and be used in modules/functions).
+        """
+
+        @torch.jit.script
+        def identity(x: torch.Tensor):
+            if isinstance(neural_type(), nelements.VoidType):
+                return x
+            else:
+                return x + 1
+
+        _ = identity(torch.tensor(1.0))
+
+    @pytest.mark.unit
+    def test_dummy_module_with_io_types(self):
+        """Test module with input/output types"""
+        module = torch.jit.script(DummyModuleWithIOTypes())
+        x = torch.rand(2, 10)
+        result = module(x)
+        assert result.shape == x.shape
+        assert torch.allclose(result, x)
+
+    @pytest.mark.unit
+    def test_chain_with_types(self):
+        """Test applying 2 modules consecutively with input/output types"""
+        dummy_module = torch.jit.script(DummyModuleWithIOTypes())
+        module = torch.jit.script(SimpleLinearWithTypes())
+        x = torch.zeros(2, 10)
+        result = module(dummy_module(x))
+        assert result.shape == (2, 20)
+        assert torch.allclose(result, torch.zeros_like(result))