feat(common): add build_module for configurable sub-module blocks

@@ -14,18 +14,30 @@ # limitations under the License. # pylint: disable=missing-function-docstring,missing-class-docstring +import functools+import inspect import logging import math import os-from typing import Iterable, List+from collections.abc import Mapping+from typing import Any, Callable, Iterable, List  logger = logging.getLogger(__name__)  import einops import torch import torch.nn as nn+from omegaconf import DictConfig++__all__ = ['if_exist', '_compute_softmax', 'flatten', 'build_module', 'block_binds', 'ModuleBlock', 'SquaredReLU']+++class SquaredReLU(nn.Module):+    """ReLU squared, ``relu(x) ** 2`` (So et al., Primer, 2021)."""++    def forward(self, x: torch.Tensor) -> torch.Tensor:+        return torch.square(torch.relu(x)) -__all__ = ['if_exist', '_compute_softmax', 'flatten']  activation_registry = {     "identity": nn.Identity,@@ -35,8 +47,81 @@     "swish": nn.SiLU,     "silu": nn.SiLU,     "gelu": nn.GELU,+    "squared_relu": SquaredReLU, } +ModuleBlock = Callable[..., nn.Module] | functools.partial | Mapping+++def build_module(block: ModuleBlock, **defaults: Any) -> nn.Module:+    """Instantiate a configurable sub-module of a parent module.++    Args:+        block: what builds the module:+            a class or callable returning it,+            a ``functools.partial`` of one+            (what Hydra creates for a config section with ``_target_`` and ``_partial_: true``),+            or such a config section itself, resolved through ``safe_instantiate``.+        defaults: the arguments the parent supplies, such as sizes and dropout rates;+            arguments bound in the partial or the section take precedence, by keyword or by position,+            so a section only names what deviates from the parent's defaults.++    Returns:+        The instantiated module.+    """+    if isinstance(block, Mapping):+        block = _section_partial(block)+    if isinstance(block, functools.partial):+        positional = _positional_names(block)+        keywords = {name: value for name, value in defaults.items() if name not in positional}+        keywords.update(block.keywords)+        return block.func(*block.args, **keywords)+    return block(**defaults)+++def block_binds(block: ModuleBlock | None, name: str) -> bool:+    """Whether a block section or partial fixes the argument ``name`` itself, by keyword or by position."""+    if isinstance(block, Mapping):+        if '_args_' not in block:+            return name in block+        block = _section_partial(block)+    if isinstance(block, functools.partial):+        return name in block.keywords or name in _positional_names(block)+    return False+++def _section_partial(section: Mapping) -> functools.partial:+    """The partial of a config section with ``_target_``, resolved through ``safe_instantiate``.++    An existing ``DictConfig`` is passed on as it is, so interpolations inside it still resolve against its root;+    a plain mapping becomes one.+    ``_partial_`` defaults to true without being written into the caller's section.+    """+    # imported here: nemo.core.classes.common imports this module+    from nemo.core.classes.common import safe_instantiate++    config = section if isinstance(section, DictConfig) else DictConfig(dict(section))+    return safe_instantiate(config, _partial_=config.get('_partial_', True))+++def _positional_names(block: functools.partial) -> frozenset[str]:+    """The parameter names a partial's positional arguments occupy.++    Arguments taken by a variadic ``*args`` occupy no name.+    For a target without an inspectable signature the result is empty:+    the positional arguments pass through unchanged and the parent's defaults all stay.+    """+    if not block.args:+        return frozenset()+    try:+        signature = inspect.signature(block.func)+        bound = signature.bind_partial(*block.args)+    except (TypeError, ValueError):+        return frozenset()+    return frozenset(+        name for name in bound.arguments if signature.parameters[name].kind is not inspect.Parameter.VAR_POSITIONAL+    )+  def if_exist(outfold: str, files: List[str]):     """

@@ -12,6 +12,7 @@ # 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 functools import os import string from contextlib import contextmanager@@ -21,9 +22,17 @@ import numpy as np import pytest import torch+from omegaconf import DictConfig, OmegaConf  from nemo.collections.common.parts.preprocessing.manifest import get_full_path, is_tarred_dataset-from nemo.collections.common.parts.utils import flatten, mask_sequence_tensor+from nemo.collections.common.parts.utils import (+    SquaredReLU,+    activation_registry,+    block_binds,+    build_module,+    flatten,+    mask_sequence_tensor,+)   class TestListUtils:@@ -366,3 +375,77 @@ def test_is_tarred_dataset(self):          # 3) no manifest file, treated as non-tarred dataset         assert not is_tarred_dataset("_file_1.wav", None)+++class TestBuildModule:+    @pytest.mark.unit+    def test_class_gets_the_parent_defaults(self):+        norm = build_module(torch.nn.LayerNorm, normalized_shape=8)+        assert isinstance(norm, torch.nn.LayerNorm) and norm.normalized_shape == (8,) and norm.eps == 1e-5++    @pytest.mark.unit+    def test_partial_keywords_override_the_parent_defaults(self):+        block = functools.partial(torch.nn.LayerNorm, eps=1e-6, elementwise_affine=False)+        norm = build_module(block, normalized_shape=8, eps=1e-5)+        assert norm.eps == 1e-6 and norm.weight is None and norm.normalized_shape == (8,)++    @pytest.mark.unit+    def test_hydra_section_is_instantiated_lazily(self):+        section = DictConfig({"_target_": "torch.nn.LayerNorm", "eps": 1e-6})+        norm = build_module(section, normalized_shape=8)+        assert isinstance(norm, torch.nn.LayerNorm) and norm.eps == 1e-6+        explicit = {"_target_": "torch.nn.LayerNorm", "_partial_": True, "eps": 1e-3}+        assert build_module(explicit, normalized_shape=4).eps == 1e-3++    @pytest.mark.unit+    def test_section_keeps_its_root_for_interpolations(self):+        cfg = OmegaConf.create(+            {"width": 4, "block": {"_target_": "torch.nn.LayerNorm", "normalized_shape": ["${width}"]}}+        )+        assert build_module(cfg.block).normalized_shape == (4,)+        assert "_partial_" not in cfg.block+        scalar = OmegaConf.create(+            {"width": 4, "block": {"_target_": "torch.nn.LayerNorm", "normalized_shape": "${width}"}}+        )+        assert build_module(scalar.block).normalized_shape == (4,)++    @pytest.mark.unit+    def test_section_targets_are_validated(self):+        with pytest.raises(Exception, match="unsafe target"):+            build_module({"_target_": "os.system"}, command="true")++    @pytest.mark.unit+    def test_positional_arguments_take_precedence(self):+        block = functools.partial(torch.nn.LayerNorm, 4)+        assert build_module(block, normalized_shape=8).normalized_shape == (4,)+        norm = build_module(block, eps=1e-3)+        assert norm.normalized_shape == (4,) and norm.eps == 1e-3+        section = {"_target_": "torch.nn.LayerNorm", "_args_": [4]}+        assert build_module(section, normalized_shape=8).normalized_shape == (4,)++    @pytest.mark.unit+    def test_variadic_positional_arguments_occupy_no_name(self):+        def factory(*shape: int, eps: float = 1e-5) -> torch.nn.LayerNorm:+            return torch.nn.LayerNorm(shape, eps=eps)++        block = functools.partial(factory, 2, 3)+        norm = build_module(block, eps=1e-3)+        assert norm.normalized_shape == (2, 3) and norm.eps == 1e-3+        assert not block_binds(block, "shape") and not block_binds(block, "eps")++    @pytest.mark.unit+    def test_block_binds(self):+        assert block_binds(functools.partial(torch.nn.LayerNorm, eps=1e-6), "eps")+        assert not block_binds(functools.partial(torch.nn.LayerNorm), "eps")+        assert block_binds({"_target_": "torch.nn.LayerNorm", "eps": 1e-6}, "eps")+        assert not block_binds(torch.nn.LayerNorm, "eps") and not block_binds(None, "eps")+        positional = functools.partial(torch.nn.LayerNorm, 4)+        assert block_binds(positional, "normalized_shape") and not block_binds(positional, "eps")+        section = {"_target_": "torch.nn.LayerNorm", "_args_": [4]}+        assert block_binds(section, "normalized_shape") and not block_binds(section, "eps")++    @pytest.mark.unit+    def test_squared_relu(self):+        x = torch.tensor([-2.0, 0.0, 3.0])+        torch.testing.assert_close(SquaredReLU()(x), torch.tensor([0.0, 0.0, 9.0]))+        assert activation_registry["squared_relu"] is SquaredReLU