chore: import upstream snapshot with attribution

nemo/core/classes/__init__.py

@@ -0,0 +1,35 @@
+# 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 hydra
+import lightning.pytorch
+import omegaconf
+
+from nemo.core.classes.common import (
+    FileIO,
+    Model,
+    PretrainedModelInfo,
+    Serialization,
+    Typing,
+    is_typecheck_enabled,
+    typecheck,
+)
+from nemo.core.classes.dataset import Dataset, IterableDataset
+from nemo.core.classes.exportable import Exportable, ExportFormat
+from nemo.core.classes.loss import Loss
+from nemo.core.classes.mixins import access_mixins, adapter_mixins, hf_io_mixin
+from nemo.core.classes.modelPT import ModelPT
+from nemo.core.classes.module import NeuralModule
+from nemo.utils import exceptions

nemo/core/classes/dataset.py

@@ -0,0 +1,108 @@
+# 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 Optional
+
+from torch.utils import data
+
+from nemo.core.classes import Serialization, Typing, typecheck
+
+__all__ = ['Dataset', 'IterableDataset']
+
+
+class Dataset(data.Dataset, Typing, Serialization):
+    """Dataset with output ports
+
+    Please Note: Subclasses of IterableDataset should *not* implement input_types.
+    """
+
+    def _collate_fn(self, batch):
+        """
+        A default implementation of a collation function.
+        Users should override this method to define custom data loaders.
+        """
+        return data.dataloader.default_collate(batch)
+
+    @typecheck()
+    def collate_fn(self, batch):
+        """
+        This is the method that user pass as functor to DataLoader.
+        The method optionally performs neural type checking and add types to the outputs.
+
+        Please note, subclasses of Dataset should not implement `input_types`.
+
+        Usage:
+
+        .. code-block:: python
+
+            dataloader = torch.utils.data.DataLoader(
+                    ....,
+                    collate_fn=dataset.collate_fn,
+                    ....
+            )
+
+        Returns:
+            Collated batch, with or without types.
+        """
+        if self.input_types is not None:
+            raise TypeError("Datasets should not implement `input_types` as they are not checked")
+
+        # Simply forward the inner `_collate_fn`
+        return self._collate_fn(batch)
+
+
+class IterableDataset(data.IterableDataset, Typing, Serialization):
+    """Iterable Dataset with output ports
+
+    Please Note: Subclasses of IterableDataset should *not* implement input_types.
+    """
+
+    def _collate_fn(self, batch):
+        """
+        A default implementation of a collation function.
+        Users should override this method to define custom data loaders.
+        """
+        return data.dataloader.default_collate(batch)
+
+    @typecheck()
+    def collate_fn(self, batch):
+        """
+        This is the method that user pass as functor to DataLoader.
+        The method optionally performs neural type checking and add types to the outputs.
+
+        # Usage:
+        dataloader = torch.utils.data.DataLoader(
+                ....,
+                collate_fn=dataset.collate_fn,
+                ....
+        )
+
+        Returns:
+            Collated batch, with or without types.
+        """
+        if self.input_types is not None:
+            raise TypeError("Datasets should not implement `input_types` as they are not checked")
+
+        # Simply forward the inner `_collate_fn`
+        return self._collate_fn(batch)
+
+
+@dataclass
+class DatasetConfig:
+    # ...
+    batch_size: int = 32
+    drop_last: bool = False
+    shuffle: bool = False
+    num_workers: Optional[int] = 0
+    pin_memory: bool = True

nemo/core/classes/exportable.py

@@ -0,0 +1,370 @@
+# 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 abc import ABC
+from typing import Dict, List, Optional, Union
+
+import torch
+from lightning.pytorch.core.module import _jit_is_scripting
+
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import NeuralType
+from nemo.core.utils.neural_type_utils import get_dynamic_axes, get_io_names
+from nemo.utils import logging, monkeypatched
+from nemo.utils.export_utils import (
+    ExportFormat,
+    augment_filename,
+    get_export_format,
+    parse_input_example,
+    rename_onnx_io,
+    replace_for_export,
+    verify_runtime,
+    verify_torchscript,
+    wrap_forward_method,
+)
+
+__all__ = ['ExportFormat', 'Exportable']
+
+
+class Exportable(ABC):
+    """
+    This Interface should be implemented by particular classes derived from nemo.core.NeuralModule or nemo.core.ModelPT.
+    It gives these entities ability to be exported for deployment to formats such as ONNX.
+
+    Usage:
+        # exporting pre-trained model to ONNX file for deployment.
+        model.eval()
+        model.to('cuda')  # or to('cpu') if you don't have GPU
+
+        model.export('mymodel.onnx', [options])  # all arguments apart from `output` are optional.
+    """
+
+    @property
+    def input_module(self):
+        return self
+
+    @property
+    def output_module(self):
+        return self
+
+    def export(
+        self,
+        output: str,
+        input_example=None,
+        verbose=False,
+        do_constant_folding=True,
+        onnx_opset_version=None,
+        check_trace: Union[bool, List[torch.Tensor]] = False,
+        dynamic_axes=None,
+        check_tolerance=0.01,
+        export_modules_as_functions=False,
+        keep_initializers_as_inputs=None,
+        use_dynamo=False,
+    ):
+        """
+        Exports the model to the specified format. The format is inferred from the file extension of the output file.
+
+        Args:
+            output (str): Output file name. File extension be .onnx, .pt, or .ts, and is used to select export
+                path of the model.
+            input_example (list or dict): Example input to the model's forward function. This is used to
+                trace the model and export it to ONNX/TorchScript. If the model takes multiple inputs, then input_example
+                should be a list of input examples. If the model takes named inputs, then input_example
+                should be a dictionary of input examples.
+            verbose (bool): If True, will print out a detailed description of the model's export steps, along with
+                the internal trace logs of the export process.
+            do_constant_folding (bool): If True, will execute constant folding optimization on the model's graph
+                before exporting. This is ONNX specific.
+            onnx_opset_version (int): The ONNX opset version to export the model to. If None, will use a reasonable
+                default version.
+            check_trace (bool): If True, will verify that the model's output matches the output of the traced
+                model, upto some tolerance.
+            dynamic_axes (dict): A dictionary mapping input and output names to their dynamic axes. This is
+                used to specify the dynamic axes of the model's inputs and outputs. If the model takes multiple inputs,
+                then dynamic_axes should be a list of dictionaries. If the model takes named inputs, then dynamic_axes
+                should be a dictionary of dictionaries. If None, will use the dynamic axes of the input_example
+                derived from the NeuralType of the input and output of the model.
+            check_tolerance (float): The tolerance to use when checking the model's output against the traced
+                model's output. This is only used if check_trace is True. Note the high tolerance is used because
+                the traced model is not guaranteed to be 100% accurate.
+            export_modules_as_functions (bool): If True, will export the model's submodules as functions. This is
+                ONNX specific.
+            keep_initializers_as_inputs (bool): If True, will keep the model's initializers as inputs in the onnx graph.
+                This is ONNX specific.
+            use_dynamo (bool): If True, use onnx.dynamo_export() instead of onnx.export(). This is ONNX specific.
+
+        Returns:
+            A tuple of two outputs.
+            Item 0 in the output is a list of outputs, the outputs of each subnet exported.
+            Item 1 in the output is a list of string descriptions. The description of each subnet exported can be
+            used for logging purposes.
+        """
+        all_out = []
+        all_descr = []
+        for subnet_name in self.list_export_subnets():
+            model = self.get_export_subnet(subnet_name)
+            out_name = augment_filename(output, subnet_name)
+            out, descr, out_example = model._export(
+                out_name,
+                input_example=input_example,
+                verbose=verbose,
+                do_constant_folding=do_constant_folding,
+                onnx_opset_version=onnx_opset_version,
+                check_trace=check_trace,
+                dynamic_axes=dynamic_axes,
+                check_tolerance=check_tolerance,
+                export_modules_as_functions=export_modules_as_functions,
+                keep_initializers_as_inputs=keep_initializers_as_inputs,
+                use_dynamo=use_dynamo,
+            )
+            # Propagate input example (default scenario, may need to be overriden)
+            if input_example is not None:
+                input_example = out_example
+            all_out.append(out)
+            all_descr.append(descr)
+            logging.info("Successfully exported {} to {}".format(model.__class__.__name__, out_name))
+        return (all_out, all_descr)
+
+    def _export(
+        self,
+        output: str,
+        input_example=None,
+        verbose=False,
+        do_constant_folding=True,
+        onnx_opset_version=None,
+        check_trace: Union[bool, List[torch.Tensor]] = False,
+        dynamic_axes=None,
+        check_tolerance=0.01,
+        export_modules_as_functions=False,
+        keep_initializers_as_inputs=None,
+        use_dynamo=False,
+    ):
+        my_args = locals().copy()
+        my_args.pop('self')
+
+        self.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+        exportables = []
+        for m in self.modules():
+            if isinstance(m, Exportable):
+                exportables.append(m)
+
+        qual_name = self.__module__ + '.' + self.__class__.__qualname__
+        format = get_export_format(output)
+        output_descr = f"{qual_name} exported to {format}"
+
+        # Pytorch's default opset version is too low, using reasonable latest one
+        if onnx_opset_version is None:
+            onnx_opset_version = 17
+
+        try:
+            # Disable typechecks
+            typecheck.set_typecheck_enabled(enabled=False)
+
+            # Allow user to completely override forward method to export
+            forward_method, old_forward_method = wrap_forward_method(self)
+
+            # Set module mode
+            with torch.inference_mode(), torch.no_grad(), torch.jit.optimized_execution(True), _jit_is_scripting():
+
+                if input_example is None:
+                    input_example = self.input_module.input_example()
+
+                # Remove i/o examples from args we propagate to enclosed Exportables
+                my_args.pop('output')
+                my_args.pop('input_example')
+
+                # Run (posibly overridden) prepare methods before calling forward()
+                for ex in exportables:
+                    ex._prepare_for_export(**my_args, noreplace=True)
+                self._prepare_for_export(output=output, input_example=input_example, **my_args)
+
+                input_list, input_dict = parse_input_example(input_example)
+                input_names = self.input_names
+                output_names = self.output_names
+                output_example = self.forward(*input_list, **input_dict)
+                if not isinstance(output_example, tuple):
+                    output_example = (output_example,)
+
+                if check_trace:
+                    if isinstance(check_trace, bool):
+                        check_trace_input = [input_example]
+                    else:
+                        check_trace_input = check_trace
+
+                if format == ExportFormat.TORCHSCRIPT:
+                    jitted_model = torch.jit.trace_module(
+                        self,
+                        {"forward": tuple(input_list) + tuple(input_dict.values())},
+                        strict=True,
+                        check_trace=check_trace,
+                        check_tolerance=check_tolerance,
+                    )
+                    jitted_model = torch.jit.freeze(jitted_model)
+                    if verbose:
+                        logging.info(f"JIT code:\n{jitted_model.code}")
+                    jitted_model.save(output)
+                    jitted_model = torch.jit.load(output)
+
+                    if check_trace:
+                        verify_torchscript(jitted_model, output, check_trace_input, check_tolerance)
+                elif format == ExportFormat.ONNX:
+                    # dynamic axis is a mapping from input/output_name => list of "dynamic" indices
+                    if dynamic_axes is None:
+                        dynamic_axes = self.dynamic_shapes_for_export(use_dynamo)
+                    # When using dynamo export, use dynamic_shapes instead of dynamic_axes
+                    if use_dynamo:
+                        dynamic_shapes = dynamic_axes
+                    if use_dynamo:
+                        typecheck.enable_wrapping(enabled=False)
+                        # https://github.com/pytorch/pytorch/issues/126339
+                        with monkeypatched(torch.nn.RNNBase, "flatten_parameters", lambda *args: None):
+                            logging.info(f"Running export.export, dynamic shapes:{dynamic_shapes}\n")
+
+                            # We have to use different types of arguments for dynamo_export to achieve
+                            # same external weights behaviour as onnx.export :
+                            # https://github.com/pytorch/pytorch/issues/126479
+                            # https://github.com/pytorch/pytorch/issues/126269
+                            mem_params = sum([param.nelement() * param.element_size() for param in self.parameters()])
+                            mem_bufs = sum([buf.nelement() * buf.element_size() for buf in self.buffers()])
+                            mem = mem_params + mem_bufs
+
+                            if mem > 2 * 1000 * 1000 * 1000:
+                                ex_model = torch.export.export(
+                                    self,
+                                    tuple(input_list),
+                                    kwargs=input_dict,
+                                    dynamic_shapes=dynamic_shapes,
+                                    strict=False,
+                                )
+                                ex_model = ex_model.run_decompositions()
+                                model_state = ex_model.state_dict
+                            else:
+                                model_state = None
+                                ex_model = self
+
+                            options = torch.onnx.ExportOptions(dynamic_shapes=True, op_level_debug=True)
+                            ex = torch.onnx.dynamo_export(ex_model, *input_list, **input_dict, export_options=options)
+                            ex.save(output, model_state=model_state)
+
+                            del ex
+                            del ex_model
+                            # Rename I/O after save - don't want to risk modifying ex._model_proto
+                            rename_onnx_io(output, input_names, output_names)
+                    else:
+                        torch.onnx.export(
+                            self,
+                            input_example,
+                            output,
+                            input_names=input_names,
+                            output_names=output_names,
+                            verbose=verbose,
+                            do_constant_folding=do_constant_folding,
+                            dynamic_axes=dynamic_axes,
+                            opset_version=onnx_opset_version,
+                            keep_initializers_as_inputs=keep_initializers_as_inputs,
+                            export_modules_as_functions=export_modules_as_functions,
+                            dynamo=False,  # Use legacy TorchScript-based exporter for LSTM compatibility
+                        )
+
+                    if check_trace:
+                        verify_runtime(self, output, check_trace_input, input_names, check_tolerance=check_tolerance)
+                else:
+                    raise ValueError(f'Encountered unknown export format {format}.')
+        finally:
+            typecheck.enable_wrapping(enabled=True)
+            typecheck.set_typecheck_enabled(enabled=True)
+            if forward_method:
+                type(self).forward = old_forward_method
+            self._export_teardown()
+        return (output, output_descr, output_example)
+
+    @property
+    def disabled_deployment_input_names(self) -> List[str]:
+        """Implement this method to return a set of input names disabled for export"""
+        return []
+
+    @property
+    def disabled_deployment_output_names(self) -> List[str]:
+        """Implement this method to return a set of output names disabled for export"""
+        return []
+
+    @property
+    def supported_export_formats(self) -> List[ExportFormat]:
+        """Implement this method to return a set of export formats supported. Default is all types."""
+        return [ExportFormat.ONNX, ExportFormat.TORCHSCRIPT]
+
+    def _prepare_for_export(self, **kwargs):
+        """
+        Override this method to prepare module for export. This is in-place operation.
+        Base version does common necessary module replacements (Apex etc)
+        """
+        if not 'noreplace' in kwargs:
+            replace_for_export(self)
+
+    def _export_teardown(self):
+        """
+        Override this method for any teardown code after export.
+        """
+        pass
+
+    @property
+    def input_names(self):
+        return get_io_names(self.input_module.input_types_for_export, self.disabled_deployment_input_names)
+
+    @property
+    def output_names(self):
+        return get_io_names(self.output_module.output_types_for_export, self.disabled_deployment_output_names)
+
+    @property
+    def input_types_for_export(self) -> Optional[Dict[str, NeuralType]]:
+        return self.input_types
+
+    @property
+    def output_types_for_export(self):
+        return self.output_types
+
+    def dynamic_shapes_for_export(self, use_dynamo=False):
+        return get_dynamic_axes(self.input_module.input_types_for_export, self.input_names, use_dynamo)
+
+    def get_export_subnet(self, subnet=None):
+        """
+        Returns Exportable subnet model/module to export
+        """
+        if subnet is None or subnet == 'self':
+            return self
+        else:
+            return getattr(self, subnet)
+
+    def list_export_subnets(self):
+        """
+        Returns default set of subnet names exported for this model
+        First goes the one receiving input (input_example)
+        """
+        return ['self']
+
+    def get_export_config(self):
+        """
+        Returns export_config dictionary
+        """
+        return getattr(self, 'export_config', {})
+
+    def set_export_config(self, args):
+        """
+        Sets/updates export_config dictionary
+        """
+        ex_config = self.get_export_config()
+        ex_config.update(args)
+        self.export_config = ex_config

nemo/core/classes/loss.py

@@ -0,0 +1,26 @@
+# 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 torch
+
+from nemo.core.classes.common import Serialization, Typing
+
+__all__ = ['Loss']
+
+
+class Loss(torch.nn.modules.loss._Loss, Typing, Serialization):
+    """Inherit this class to implement custom loss."""
+
+    def __init__(self, **kwargs):
+        super(Loss, self).__init__(**kwargs)

nemo/core/classes/mixins/__init__.py

@@ -0,0 +1,28 @@
+# 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.
+
+from nemo.core.classes.mixins.access_mixins import AccessMixin, set_access_cfg
+from nemo.core.classes.mixins.adapter_mixin_strategies import (
+    ResidualAddAdapterStrategy,
+    ResidualAddAdapterStrategyConfig,
+    ReturnResultAdapterStrategy,
+    ReturnResultAdapterStrategyConfig,
+)
+from nemo.core.classes.mixins.adapter_mixins import (
+    AdapterModelPTMixin,
+    AdapterModuleMixin,
+    get_registered_adapter,
+    register_adapter,
+)
+from nemo.core.classes.mixins.hf_io_mixin import HuggingFaceFileIO

nemo/core/classes/mixins/access_mixins.py

@@ -0,0 +1,146 @@
+# 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 abc import ABC
+from typing import Optional
+
+import torch
+from omegaconf import DictConfig
+
+_DEFAULT_ACCESS_GUID = "default"
+_ACCESS_CFG = DictConfig({_DEFAULT_ACCESS_GUID: {"detach": False, "convert_to_cpu": False}})
+_ACCESS_ENABLED = {_DEFAULT_ACCESS_GUID: False}
+
+
+def set_access_cfg(cfg: DictConfig, guid: Optional[str] = None):
+    if cfg is None or not isinstance(cfg, DictConfig):
+        raise TypeError(f"cfg must be a DictConfig")
+    global _ACCESS_CFG
+    global _DEFAULT_ACCESS_GUID
+    if guid is not None:
+        _ACCESS_CFG[guid] = cfg
+    else:
+        _ACCESS_CFG[_DEFAULT_ACCESS_GUID] = cfg
+
+
+class AccessMixin(ABC):
+    """
+    Allows access to output of intermediate layers of a model
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._registry = {}  # dictionary of lists
+
+    def register_accessible_tensor(self, name, tensor):
+        """
+        Register tensor for later use.
+        """
+        if self.access_cfg.get('convert_to_cpu', False):
+            tensor = tensor.cpu()
+
+        if self.access_cfg.get('detach', False):
+            tensor = tensor.detach()
+
+        if not hasattr(self, '_registry'):
+            self._registry = {}
+
+        if name not in self._registry:
+            self._registry[name] = []
+
+        self._registry[name].append(tensor)
+
+    @classmethod
+    def get_module_registry(cls, module: torch.nn.Module):
+        """
+        Extract all registries from named submodules, return dictionary where
+        the keys are the flattened module names, the values are the internal registry
+        of each such module.
+        """
+        module_registry = {}
+        for name, m in module.named_modules():
+            if hasattr(m, '_registry') and len(m._registry) > 0:
+                module_registry[name] = m._registry
+        return module_registry
+
+    def reset_registry(self: torch.nn.Module, registry_key: Optional[str] = None):
+        """
+        Reset the registries of all named sub-modules
+        """
+        if hasattr(self, "_registry"):
+            if registry_key is None:
+                self._registry.clear()
+            else:
+                if registry_key in self._registry:
+                    self._registry.pop(registry_key)
+                else:
+                    raise KeyError(
+                        f"Registry key `{registry_key}` provided, but registry does not have this key.\n"
+                        f"Available keys in registry : {list(self._registry.keys())}"
+                    )
+
+        for _, m in self.named_modules():
+            if hasattr(m, "_registry"):
+                if registry_key is None:
+                    m._registry.clear()
+                else:
+                    if registry_key in self._registry:
+                        self._registry.pop(registry_key)
+                    else:
+                        raise KeyError(
+                            f"Registry key `{registry_key}` provided, but registry does not have this key.\n"
+                            f"Available keys in registry : {list(self._registry.keys())}"
+                        )
+
+        # Explicitly disable registry cache after reset
+        AccessMixin.set_access_enabled(access_enabled=False, guid=getattr(self, "model_guid", None))
+
+    @property
+    def access_cfg(self):
+        """
+        Returns:
+            The global access config shared across all access mixin modules.
+        """
+        global _ACCESS_CFG
+        global _DEFAULT_ACCESS_GUID
+        guid = self.model_guid if getattr(self, "model_guid", None) else _DEFAULT_ACCESS_GUID
+        if hasattr(self, "propagate_model_guid"):
+            self.propagate_model_guid()
+        if guid not in _ACCESS_CFG:
+            _ACCESS_CFG[guid] = DictConfig({})
+        return _ACCESS_CFG[guid]
+
+    @classmethod
+    def update_access_cfg(cls, cfg: dict, guid: Optional[str] = None):
+        global _ACCESS_CFG
+        global _DEFAULT_ACCESS_GUID
+        guid = guid if guid is not None else _DEFAULT_ACCESS_GUID
+        if guid not in _ACCESS_CFG:
+            _ACCESS_CFG[guid] = cfg
+        else:
+            _ACCESS_CFG[guid].update(cfg)
+
+    @classmethod
+    def is_access_enabled(cls, guid: Optional[str] = None):
+        global _ACCESS_ENABLED
+        global _DEFAULT_ACCESS_GUID
+        guid = guid if guid is not None else _DEFAULT_ACCESS_GUID
+        return _ACCESS_ENABLED.get(guid, False)
+
+    @classmethod
+    def set_access_enabled(cls, access_enabled: bool, guid: Optional[str] = None):
+        global _ACCESS_ENABLED
+        global _DEFAULT_ACCESS_GUID
+        guid = guid if guid is not None else _DEFAULT_ACCESS_GUID
+        _ACCESS_ENABLED[guid] = access_enabled

nemo/core/classes/mixins/adapter_mixin_strategies.py

@@ -0,0 +1,259 @@
+# 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.
+
+from abc import ABC
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple, Union
+
+import torch
+
+from nemo.core.classes.mixins import AccessMixin
+
+
+class AbstractAdapterStrategy(ABC):
+    def forward(self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'):
+        """
+        Forward method that defines how the output of the adapter should be merged with the input, or if it
+        should be merged at all.
+
+        Also provides the module that called this strategy - thereby allowing access to all other
+        adapters in the calling module. This can be useful if one adapter is a meta adapter, that
+        combines the outputs of various adapters. In such a case, the input can be forwarded across
+        all other adapters, collecting their outputs, and those outputs can then be merged via some
+        strategy. For example, refer to :
+
+        - [AdapterFusion: Non-Destructive Task Composition for Transfer Learning](https://arxiv.org/abs/2005.00247)
+        - [Exploiting Adapters for Cross-lingual Low-resource Speech Recognition](https://arxiv.org/abs/2105.11905)
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        raise NotImplementedError()
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+
+class ReturnResultAdapterStrategy(AbstractAdapterStrategy):
+    """
+    An implementation of an adapter strategy that simply returns the result of the adapter.
+    Supports stochastic
+    """
+
+    def forward(self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'):
+        """
+        A basic strategy, which simply returns the result of the adapter's calculation as the output.
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        result = self.compute_output(input, adapter, module=module)
+
+        return result
+
+    def compute_output(
+        self,
+        input: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor], Dict[str, Any]],
+        adapter: torch.nn.Module,
+        *,
+        module: 'AdapterModuleMixin',
+    ) -> torch.Tensor:
+        """
+        Compute the output of a single adapter to some input.
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        if isinstance(input, (list, tuple)):
+            out = adapter(*input)
+        elif isinstance(input, dict):
+            out = adapter(**input)
+        else:
+            out = adapter(input)
+        return out
+
+
+@dataclass
+class ReturnResultAdapterStrategyConfig:
+    _target_: str = "{0}.{1}".format(
+        ReturnResultAdapterStrategy.__module__, ReturnResultAdapterStrategy.__name__
+    )  # mandatory field
+
+
+class ResidualAddAdapterStrategy(AbstractAdapterStrategy):
+    """
+    An implementation of residual addition of an adapter module with its input.
+    Supports stochastic depth regularization.
+    """
+
+    def __init__(self, stochastic_depth: float = 0.0, l2_lambda: float = 0.0):
+        """
+        An implementation of residual addition of an adapter module with its input.
+        Performs output = input + adapter(input).
+
+        Args:
+            stochastic_depth: float, when greater than one, can optionally dropout the output of
+                the adapter's forward pass.
+            l2_lambda: L2 norm of the difference between the original input to the function, and the adapter's
+                output result. Disabled if set to 0.0.
+        """
+        super().__init__()
+        self.stochastic_depth = stochastic_depth
+        self.l2_lambda = l2_lambda
+
+    def forward(self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'):
+        """
+        A basic strategy, comprising of a residual connection over the input, after forward pass by
+        the underlying adapter.
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        out = self.compute_output(input, adapter, module=module)
+
+        # If not in training mode, or probability of stochastic depth is 0, skip step.
+        p = self.stochastic_depth
+        if not module.training or p == 0.0:
+            pass
+        else:
+            out = self.apply_stochastic_depth(out, input, adapter, module=module)
+
+        # Return the residual connection output = input + adapter(input)
+        result = input + out
+
+        # If l2_lambda is activated, register the loss value
+        self.compute_auxiliary_losses(result, input, adapter, module=module)
+
+        return result
+
+    def compute_output(
+        self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'
+    ) -> torch.Tensor:
+        """
+        Compute the output of a single adapter to some input.
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        out = adapter(input)
+        return out
+
+    def apply_stochastic_depth(
+        self, output: torch.Tensor, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'
+    ):
+        """
+        Compute and apply stochastic depth if probability is greater than 0.
+
+        Args:
+            output: The result tensor, after one of the active adapters has finished its forward passes.
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after stochastic depth has been potentially applied to it.
+        """
+        # Perform stochastic depth if needed.
+        p = self.stochastic_depth
+        if p < 0.0 or p > 1.0:
+            raise ValueError(f"Stochastic depth probability has to be between 0 and 1, but got {p}")
+
+        # Apply stochastic depth to the output of adapter.
+        keep_prob = 1.0 - p
+        shape = [1] * output.ndim
+        noise = torch.empty(shape, dtype=output.dtype, device=output.device)
+        noise = noise.bernoulli_(keep_prob)
+        if keep_prob > 0.0:  # Done to normalize activation for inference mode
+            noise.div_(keep_prob)
+
+        output = noise * output
+
+        return output
+
+    def compute_auxiliary_losses(
+        self, output: torch.Tensor, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'
+    ):
+        """
+        Compute any auxiliary losses and preserve it in the tensor registry.
+
+        Args:
+            output: The result tensor, after one of the active adapters has finished its forward passes.
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+        """
+        if module.training and self.l2_lambda > 0.0:
+            if not isinstance(adapter, AccessMixin):
+                raise ValueError(f"Module {adapter.__class__.__name__} does not implement AccessMixin !")
+
+            # Only add auxiliary loss if adapter has trainable parameters that require gradients
+            if next(adapter.parameters()).requires_grad is True:
+                # Check if globally allowed to compute aux loss
+                compute_aux_loss = adapter.access_cfg.get('compute_adapter_loss', True)
+
+                if compute_aux_loss:
+                    # if l2 lambda is enabled, also enable AccessMixin
+                    adapter.set_access_enabled(access_enabled=True, guid=getattr(self, "model_guid", None))
+
+                    l2_loss = self.l2_lambda * (input - output).square().reshape(input.size(0), -1).sum(dim=-1).mean()
+                    adapter.register_accessible_tensor(name='adapter_loss', tensor=l2_loss)
+
+
+@dataclass
+class ResidualAddAdapterStrategyConfig:
+    stochastic_depth: float = 0.0
+    l2_lambda: float = 0.0
+
+    _target_: str = "{0}.{1}".format(
+        ResidualAddAdapterStrategy.__module__, ResidualAddAdapterStrategy.__name__
+    )  # mandatory field

nemo/core/classes/mixins/hf_io_mixin.py

@@ -0,0 +1,264 @@
+# 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.
+
+from abc import ABC
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+from huggingface_hub import HfApi, ModelCard, ModelCardData
+from huggingface_hub import get_token as get_hf_token
+from huggingface_hub.hf_api import ModelInfo
+from huggingface_hub.utils import SoftTemporaryDirectory
+
+
+class HuggingFaceFileIO(ABC):
+    """
+    Mixin that provides Hugging Face file IO functionality for NeMo models.
+    It is usually implemented as a mixin to `ModelPT`.
+
+    This mixin provides the following functionality:
+    - `search_huggingface_models()`: Search the hub programmatically via some model filter.
+    - `push_to_hf_hub()`: Push a model to the hub.
+    """
+
+    @classmethod
+    def get_hf_model_filter(cls) -> Dict[str, Any]:
+        """
+        Generates a filter for HuggingFace models.
+
+        Additionaly includes default values of some metadata about results returned by the Hub.
+
+        Metadata:
+            resolve_card_info: Bool flag, if set, returns the model card metadata. Default: False.
+            limit_results: Optional int, limits the number of results returned.
+
+        Returns:
+            A dict representing the arguments passable to huggingface list_models().
+        """
+        model_filter = dict(
+            author=None,
+            filter=['nemo'],
+            model_name=None,
+            limit=None,
+            full=None,
+            cardData=False,
+        )
+
+        return model_filter
+
+    @classmethod
+    def search_huggingface_models(cls, model_filter: Optional[Dict[str, Any]] = None) -> Iterable['ModelInfo']:
+        """
+        Should list all pre-trained models available via Hugging Face Hub.
+
+        The following metadata can be passed via the `model_filter` for additional results.
+        Metadata:
+
+            resolve_card_info: Bool flag, if set, returns the model card metadata. Default: False.
+
+            limit_results: Optional int, limits the number of results returned.
+
+        .. code-block:: python
+
+            # You can replace <DomainSubclass> with any subclass of ModelPT.
+            from nemo.core import ModelPT
+
+            # Get default filter dict
+            filt = <DomainSubclass>.get_hf_model_filter()
+
+            # Make any modifications to the filter as necessary
+            filt['filter'].append('en')  # Add language filter
+            filt['filter'].append('automatic-speech-recognition')  # Add task filter
+
+            # Add any metadata to the filter as needed (kwargs to list_models)
+            filt['limit'] = 5
+
+            # Obtain model info
+            model_infos = <DomainSubclass>.search_huggingface_models(model_filter=filt)
+
+            # Browse through cards and select an appropriate one
+            card = model_infos[0]
+
+            # Restore model using `modelId` of the card.
+            model = ModelPT.from_pretrained(card.modelId)
+
+        Args:
+            model_filter: Optional Dictionary (for Hugging Face Hub kwargs)
+                that filters the returned list of compatible model cards, and selects all results from each filter.
+                Users can then use `model_card.modelId` in `from_pretrained()` to restore a NeMo Model.
+
+        Returns:
+            A list of ModelInfo entries.
+        """
+        # Resolve model filter if not provided as argument
+        if model_filter is None:
+            model_filter = cls.get_hf_model_filter()
+
+        # Check if api token exists, use if it does
+        hf_token = get_hf_token()
+
+        # Search for all valid models after filtering
+        api = HfApi()
+
+        results = api.list_models(token=hf_token, sort="lastModified", **model_filter)  # type: Iterable[ModelInfo]
+
+        return results
+
+    def push_to_hf_hub(
+        self,
+        repo_id: str,
+        *,
+        pack_nemo_file: bool = True,
+        model_card: Optional['ModelCard'] | object | str = None,
+        commit_message: str = "Push model using huggingface_hub.",
+        private: bool = False,
+        api_endpoint: Optional[str] = None,
+        token: Optional[str] = None,
+        branch: Optional[str] = None,
+        allow_patterns: Optional[Union[List[str], str]] = None,
+        ignore_patterns: Optional[Union[List[str], str]] = None,
+        delete_patterns: Optional[Union[List[str], str]] = None,
+    ):
+        """
+        Upload model checkpoint to the Hub.
+
+        Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
+        `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more
+        details.
+
+        Args:
+            repo_id (`str`):
+                ID of the repository to push to (example: `"username/my-model"`).
+            pack_nemo_file (`bool`, *optional*, defaults to `True`): Whether to pack the model checkpoint and
+                configuration into a single `.nemo` file. If set to false, uploads the contents of the directory
+                containing the model checkpoint and configuration plus additional artifacts.
+            model_card (`ModelCard`, *optional*): Model card to upload with the model. If None, will use the model
+                card template provided by the class itself via `generate_model_card()`. Any object that implements
+                str(obj) can be passed here. Two keyword replacements are passed to `generate_model_card()`:
+                `model_name` and `repo_id`. If the model card generates a string, and it contains `{model_name}` or
+                `{repo_id}`, they will be replaced with the actual values.
+            commit_message (`str`, *optional*):
+                Message to commit while pushing.
+            private (`bool`, *optional*, defaults to `False`):
+                Whether the repository created should be private.
+            api_endpoint (`str`, *optional*):
+                The API endpoint to use when pushing the model to the hub.
+            token (`str`, *optional*):
+                The token to use as HTTP bearer authorization for remote files. By default, it will use the token
+                cached when running `huggingface-cli login`.
+            branch (`str`, *optional*):
+                The git branch on which to push the model. This defaults to `"main"`.
+            allow_patterns (`List[str]` or `str`, *optional*):
+                If provided, only files matching at least one pattern are pushed.
+            ignore_patterns (`List[str]` or `str`, *optional*):
+                If provided, files matching any of the patterns are not pushed.
+            delete_patterns (`List[str]` or `str`, *optional*):
+                If provided, remote files matching any of the patterns will be deleted from the repo.
+
+        Returns:
+            The url of the uploaded HF repo.
+        """
+        if "/" not in repo_id or len(repo_id.split("/")) != 2:
+            raise ValueError("Invalid repo_id provided. Please provide a repo_id of the form `username/repo-name`.")
+
+        domain_name, model_name = repo_id.split("/")
+
+        if token is None:
+            token = get_hf_token()
+
+        api = HfApi(endpoint=api_endpoint, token=token)
+        repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
+
+        # Push the files to the repo in a single commit
+        with SoftTemporaryDirectory() as tmp:
+            saved_path = Path(tmp) / repo_id
+            saved_path.mkdir(parents=True, exist_ok=True)
+
+            # Save nemo file in temp dir
+            # Get SaveRestoreConnector from subclass implementation
+            if not hasattr(self, '_save_restore_connector'):
+                raise NotImplementedError(
+                    "Model must implement a `_save_restore_connector` property to push to the HuggingFace Hub."
+                )
+
+            # We want to save a NeMo file, but not pack its contents into a tarfile by default
+            save_restore_connector = self._save_restore_connector
+            save_restore_connector.pack_nemo_file = pack_nemo_file
+
+            nemo_filepath = saved_path / f"{model_name}.nemo"
+            self.save_to(nemo_filepath)
+
+            # Save model card in temp dir
+            if model_card is None:
+                card_model_name = model_name.replace("_", " ").split(" ")
+                card_model_name = " ".join([word.capitalize() for word in card_model_name])
+                template_kwargs = {
+                    'model_name': card_model_name,
+                    'repo_id': repo_id,
+                }
+                # Generate model card from subclass that implements this method
+                model_card = self.generate_model_card(type='hf', template_kwargs=template_kwargs)
+
+            # Convert model card to str
+            model_card = str(model_card)
+
+            # Write model card to temp dir
+            model_card_filepath = saved_path / "README.md"
+            model_card_filepath.write_text(str(model_card), encoding='utf-8', errors='ignore')
+
+            api.upload_folder(
+                repo_id=repo_id,
+                repo_type="model",
+                folder_path=saved_path,
+                commit_message=commit_message,
+                revision=branch,
+                allow_patterns=allow_patterns,
+                ignore_patterns=ignore_patterns,
+                delete_patterns=delete_patterns,
+            )
+
+            if branch is None:
+                branch = "main"
+
+            return f"https://huggingface.co/{repo_id}/tree/{branch}"
+
+    def _get_hf_model_card(self, template: str, template_kwargs: Optional[Dict[str, str]] = None):
+        """
+        Generate a HuggingFace ModelCard from a str template. The template may have markers with `{key}` that will be
+        populated by values from `template_kwargs` if provided.
+
+        Args:
+            template: Str template for the model card.
+            template_kwargs (optional): Dict of key-value pairs to populate the template with.
+
+        Returns:
+            A HuggingFace ModelCard object that can be converted to a model card string.
+        """
+        card_data = ModelCardData(
+            library_name='nemo',
+            tags=['pytorch', 'NeMo'],
+            license='cc-by-4.0',
+            ignore_metadata_errors=True,
+        )
+
+        if 'card_data' not in template_kwargs:
+            template_kwargs['card_data'] = card_data.to_yaml()
+
+        # Update template with kwargs
+        # We need to do a manual replace because not all keys may be provided in the kwargs
+        for key, val in template_kwargs.items():
+            template = template.replace("{" + key.strip() + "}", val)
+
+        hf_model_card = ModelCard(template)
+        return hf_model_card

nemo/core/classes/module.py

@@ -0,0 +1,141 @@
+# 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 contextlib import contextmanager
+
+import torch
+from torch.nn import Module
+
+from nemo.core.classes.common import FileIO, Serialization, Typing
+from nemo.utils import logging
+
+__all__ = ['NeuralModule', 'freeze', 'unfreeze']
+
+
+def freeze(module: Module) -> None:
+    """Freeze all parameters of ``module`` and snapshot their prior ``requires_grad`` state.
+
+    The snapshot is stored on ``module._frozen_grad_map`` so a later call to ``unfreeze(..., partial=True)``
+    can restore the pre-freeze state instead of unconditionally enabling gradients.
+    """
+    grad_map = {pname: param.requires_grad for pname, param in module.named_parameters()}
+    for param in module.parameters():
+        param.requires_grad = False
+    if not hasattr(module, '_frozen_grad_map'):
+        module._frozen_grad_map = grad_map
+    else:
+        module._frozen_grad_map.update(grad_map)
+    module.eval()
+
+
+def unfreeze(module: Module, partial: bool = False) -> None:
+    """Unfreeze parameters of ``module``.
+
+    If ``partial=True``, restore each parameter's ``requires_grad`` from the snapshot recorded by
+    ``freeze(module)``; otherwise enable gradients on every parameter. The snapshot is cleared in
+    both cases and ``module.train()`` is called.
+    """
+    if partial and not hasattr(module, '_frozen_grad_map'):
+        raise ValueError("Cannot unfreeze partially without first freezing the module with `freeze()`")
+
+    for pname, param in module.named_parameters():
+        if not partial:
+            param.requires_grad = True
+        elif pname in module._frozen_grad_map:
+            param.requires_grad = module._frozen_grad_map[pname]
+        else:
+            logging.warning(
+                f"Parameter {pname} not found in list of previously frozen parameters. Unfreezing this parameter."
+            )
+            param.requires_grad = True
+
+    if hasattr(module, '_frozen_grad_map'):
+        delattr(module, '_frozen_grad_map')
+
+    module.train()
+
+
+class NeuralModule(Module, Typing, Serialization, FileIO):
+    """
+    Abstract class offering interface shared between all PyTorch Neural Modules.
+    """
+
+    @property
+    def num_weights(self):
+        """
+        Utility property that returns the total number of parameters of NeuralModule.
+        """
+        return self._num_weights()
+
+    @torch.jit.ignore
+    def _num_weights(self):
+        num: int = 0
+        for p in self.parameters():
+            if p.requires_grad:
+                num += p.numel()
+        return num
+
+    def input_example(self, max_batch=None, max_dim=None):
+        """
+        Override this method if random inputs won't work
+        Returns:
+            A tuple sample of valid input data.
+        """
+
+        return None
+
+    def freeze(self) -> None:
+        r"""Freeze all params for inference. See :func:`freeze` for details."""
+        freeze(self)
+
+    def unfreeze(self, partial: bool = False) -> None:
+        """Unfreeze parameters for training. See :func:`unfreeze` for details.
+
+        Example:
+            ```python
+            model.encoder.freeze()              # caller freezes encoder
+            model.freeze()                      # freezes everything; encoder snapshot preserved
+            model.unfreeze(partial=True)        # decoder unfrozen, encoder stays frozen
+            ```
+        """
+        unfreeze(self, partial=partial)
+
+    @contextmanager
+    def as_frozen(self):
+        """
+        Context manager which temporarily freezes a module, yields control and finally unfreezes the module partially
+        to return to original state.
+
+        Allows for either total unfreeze or partial unfreeze (if the module was explicitly frozen
+        previously with `freeze()`). The `partial` argument is used to determine whether to unfreeze
+        all parameters or only the parameters that were previously unfrozen prior `freeze()`.
+
+        Example:
+            with model.as_frozen():  # by default, partial = True
+                # Do something with the model
+                pass
+
+            # Model's parameters are now back to original state of requires_grad
+        """
+        training_mode = self.training
+        self.freeze()
+        try:
+            yield
+        finally:
+            self.unfreeze(partial=True)
+
+            if training_mode:
+                self.train()
+            else:
+                self.eval()