returnn 1.20240626.135241__tar.gz → 1.20240626.231046__tar.gz

{returnn-1.20240626.135241 → returnn-1.20240626.231046}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240626.135241
+Version: 1.20240626.231046
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20240626.231046/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240626.231046'
+long_version = '1.20240626.231046+git.14f0203'

{returnn-1.20240626.135241 → returnn-1.20240626.231046}/returnn/datasets/distrib_files.py

@@ -159,6 +159,7 @@ class DistributeFilesDataset(CachedDataset2):
         self._file_sizes: Optional[Dict[str, int]] = None  # key -> size. for equal distribution across sub epochs
         self._data_keys: Optional[List[str]] = None
         self._num_seqs: Optional[int] = None
+        self.distrib_shard_files = distrib_shard_files
         self._shard_index, self._num_shards = _get_rank_and_size() if distrib_shard_files else (0, 1)
 
         self._workers: Dict[int, _WorkerProcParent] = {}  # epoch -> worker

{returnn-1.20240626.135241 → returnn-1.20240626.231046}/returnn/frontend/build_from_dict.py

@@ -3,8 +3,10 @@ Construct modules (or other objects) from dictionaries.
 """
 
 from __future__ import annotations
-from typing import Union, Any, Type, Dict
+from types import FunctionType, BuiltinFunctionType
+from typing import Union, Any, Type, Callable, Dict
 import importlib
+import functools
 import returnn.frontend as rf
 
 
@@ -20,17 +22,23 @@ def build_from_dict(d: Dict[str, Any], *args, **kwargs) -> Union[rf.Module, Any]
     :param d: dictionary with the class name and other options
     :param args: passed to the class constructor
     :param kwargs: passed to the class constructor
-    :return: cls(*args, **d, **kwargs)
+    :return: ``cls(*args, **d, **kwargs)`` (after "class" was popped from ``d``)
     """
     if "class" not in d:
         raise ValueError(f"build_from_dict: Missing 'class' key in dict: {d}")
     d = d.copy()
     cls_name = d.pop("class")
     cls = _get_cls(cls_name)
+    if isinstance(cls, (FunctionType, BuiltinFunctionType)):
+        if args or kwargs or d:
+            return functools.partial(cls, *args, **d, **kwargs)
+        return cls
+    if not isinstance(cls, type):
+        raise ValueError(f"build_from_dict: Expected class, got {cls!r} for class_name {cls_name!r}")
     return cls(*args, **d, **kwargs)
 
 
-def build_dict(cls: Type, **kwargs) -> Dict[str, Any]:
+def build_dict(cls: Union[Type, FunctionType, BuiltinFunctionType, Callable[..., Any]], **kwargs) -> Dict[str, Any]:
     """
     Build a dictionary for :func:`build_from_dict`.
     The class name is stored in the `"class"` key.
@@ -40,7 +48,15 @@ def build_dict(cls: Type, **kwargs) -> Dict[str, Any]:
     (e.g. for Sisyphus :func:`sis_hash_helper`)
     which should not change if the class is renamed or moved
     to keep the hash stable.
+
+    :param cls: class or function. actually we do not expect any other Callable here.
+        It's just Callable in the type hint because MyPy/PyCharm don't correctly support FunctionType
+            (https://github.com/python/mypy/issues/3171).
+    :param kwargs: other kwargs put into the dict. expect to contain only other serializable values.
+    :return: build dict. can easily be serialized. to be used with :func:`build_from_dict`.
     """
+    if not isinstance(cls, (type, FunctionType, BuiltinFunctionType)):
+        raise TypeError(f"build_dict: Expected class or function, got: {cls!r}")
     return {"class": _get_cls_name(cls), **kwargs}
 
 

{returnn-1.20240626.135241 → returnn-1.20240626.231046}/returnn/frontend/encoder/conformer.py

@@ -21,7 +21,14 @@ class ConformerPositionwiseFeedForward(rf.Module):
         FF -> Activation -> Dropout -> FF
     """
 
-    def __init__(self, out_dim: Dim, *, ff_dim: Dim, dropout: float, activation: Callable[[Tensor], Tensor]):
+    def __init__(
+        self,
+        out_dim: Dim,
+        *,
+        ff_dim: Dim,
+        dropout: float,
+        activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module],
+    ):
         """
         :param out_dim: output feature dimension
         :param ff_dim: dimension of the feed-forward layers
@@ -33,6 +40,10 @@ class ConformerPositionwiseFeedForward(rf.Module):
         self.out_dim = out_dim
         self.dropout = dropout
         self.dropout_broadcast = rf.dropout_broadcast_default()
+        if isinstance(activation, dict):
+            activation = rf.build_from_dict(activation)
+        elif not callable(activation):
+            raise TypeError(f"{self}: unexpected activation type {activation!r}")
         self.activation = activation
 
         self.linear_ff = rf.Linear(out_dim, ff_dim)
@@ -67,6 +78,8 @@ class ConformerConvBlock(rf.Module):
             out_dim, out_dim, filter_size=kernel_size, groups=out_dim.dimension, padding="same"
         )
         self.positionwise_conv2 = rf.Linear(out_dim, out_dim)
+        if not callable(norm):
+            raise TypeError(f"{self}: unexpected norm type {norm!r}")
         self.norm = norm
 
     def __call__(self, inp: Tensor, *, spatial_dim: Dim) -> Tensor:
@@ -176,13 +189,13 @@ class ConformerEncoderLayer(rf.Module):
         out_dim: Dim = Dim(512, name="conformer-enc-default-out-dim"),
         *,
         ff_dim: Dim = NotSpecified,
-        ff_activation: Callable[[Tensor], Tensor] = rf.swish,
+        ff_activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module] = rf.swish,
         dropout: float = 0.1,
         conv_kernel_size: int = 32,
-        conv_norm: Union[rf.BatchNorm, type, Any] = NotSpecified,
+        conv_norm: Union[rf.BatchNorm, type, Dict[str, Any], Any] = NotSpecified,
         conv_norm_opts: Optional[Dict[str, Any]] = None,
         num_heads: int = 4,
-        self_att: Optional[Union[rf.RelPosSelfAttention, rf.Module, type, Any]] = None,
+        self_att: Optional[Union[rf.RelPosSelfAttention, rf.Module, type, Dict[str, Any], Any]] = None,
         self_att_opts: Optional[Dict[str, Any]] = None,
         att_dropout: float = 0.1,
     ):
@@ -234,7 +247,7 @@ class ConformerEncoderLayer(rf.Module):
         self.conv_block = ConformerConvBlock(out_dim=out_dim, kernel_size=conv_kernel_size, norm=conv_norm)
         self.conv_layer_norm = rf.LayerNorm(out_dim)
 
-        if self_att is None or isinstance(self_att, type):
+        if self_att is None or isinstance(self_att, (dict, type)):
             self_att_opts_ = dict(
                 in_dim=out_dim,
                 proj_dim=out_dim,
@@ -247,9 +260,16 @@ class ConformerEncoderLayer(rf.Module):
                 self_att_opts_.update(self_att_opts)
             if self_att is None:
                 self.self_att = rf.RelPosSelfAttention(**self_att_opts_)
-            else:
+            elif isinstance(self_att, type):
                 self.self_att = self_att(**self_att_opts_)
+            elif isinstance(self_att, dict):
+                self_att_opts_ = {k: v for (k, v) in self_att_opts_.items() if k not in self_att}
+                self.self_att = rf.build_from_dict(self_att, **self_att_opts_)
+            else:
+                raise TypeError(f"{self}: invalid type: self_att {self_att!r}")
         else:
+            if not callable(self_att):
+                raise TypeError(f"{self}: invalid non-callable: self_att {self_att!r}")
             self.self_att = self_att
         self.self_att_layer_norm = rf.LayerNorm(out_dim)
 
@@ -296,13 +316,13 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         input_layer: Union[ConformerConvSubsample, ISeqDownsamplingEncoder, rf.Module, Any],
         input_dropout: float = 0.1,
         ff_dim: Dim = NotSpecified,
-        ff_activation: Callable[[Tensor], Tensor] = rf.swish,
+        ff_activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module] = rf.swish,
         dropout: float = 0.1,
         conv_kernel_size: int = 32,
-        conv_norm: Union[rf.BatchNorm, type, Any] = NotSpecified,
+        conv_norm: Union[rf.BatchNorm, type, Dict[str, Any], Any] = NotSpecified,
         num_heads: int = 4,
         att_dropout: float = 0.1,
-        encoder_layer: Optional[Union[ConformerEncoderLayer, rf.Module, type, Any]] = None,
+        encoder_layer: Optional[Union[ConformerEncoderLayer, rf.Module, type, Dict[str, Any], Any]] = None,
         encoder_layer_opts: Optional[Dict[str, Any]] = None,
         sequential=rf.Sequential,
     ):
@@ -337,7 +357,7 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         )
         self.input_dropout = input_dropout
 
-        if not encoder_layer or isinstance(encoder_layer, type):
+        if not encoder_layer or isinstance(encoder_layer, (dict, type)):
             encoder_layer_opts_ = dict(
                 out_dim=out_dim,
                 ff_dim=ff_dim,
@@ -354,8 +374,21 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
                 encoder_layer = ConformerEncoderLayer(**encoder_layer_opts_)
             elif isinstance(encoder_layer, type):
                 encoder_layer = encoder_layer(**encoder_layer_opts_)
+            elif isinstance(encoder_layer, dict):
+                # Note: Reuse all the encoder_layer_opts_.
+                # If this does not make sense for the specific encoder_layer class here,
+                # we would suggest to use a different ConformerEncoder class.
+                # (The alternative, to not reuse encoder_layer_opts_ here,
+                #  would probably be more confusing, as those options are all ignored then.
+                #  It's also not clear what args to pass then and what not.)
+                # (Maybe we should do a ConformerEncoderV2 if this is confusing here...)
+                encoder_layer_opts_ = {k: v for (k, v) in encoder_layer_opts_.items() if k not in encoder_layer}
+                encoder_layer = rf.build_from_dict(encoder_layer, **encoder_layer_opts_)
             else:
                 raise TypeError(f"unexpected encoder_layer {encoder_layer!r}")
+        else:
+            if not callable(encoder_layer):
+                raise TypeError(f"{self}: invalid non-callable encoder_layer {encoder_layer!r}")
 
         self.layers = sequential(_copy.deepcopy(encoder_layer) for _ in range(num_layers))
 

{returnn-1.20240626.135241 → returnn-1.20240626.231046}/returnn/frontend/math_.py

@@ -58,6 +58,7 @@ __all__ = [
     "floor",
     "round",
     "relu",
+    "relu_square",
     "elu",
     "selu",
     "silu",
@@ -493,6 +494,13 @@ def relu(a: Tensor) -> Tensor:
     return a._raw_backend.activation(a, "relu")
 
 
+def relu_square(x: Tensor) -> Tensor:
+    """
+    :return: ``relu(x) ** 2``
+    """
+    return square(relu(x))
+
+
 def elu(a: Tensor) -> Tensor:
     """elu"""
     # noinspection PyProtectedMember

{returnn-1.20240626.135241 → returnn-1.20240626.231046}/returnn/util/basic.py

@@ -3979,6 +3979,22 @@ def close_all_fds_except(except_fds):
             os.closerange(except_fds[i] + 1, except_fds[i + 1])
 
 
+def is_valid_fd(fd: int) -> bool:
+    """
+    :return: whether the file descriptor (fd) is still open and valid
+    """
+    # https://stackoverflow.com/questions/12340695/how-to-check-if-a-given-file-descriptor-stored-in-a-variable-is-still-valid
+    import fcntl
+    import errno
+
+    try:
+        fcntl.fcntl(fd, fcntl.F_GETFD)
+        return True  # no matter what fcntl returned, it's a valid fd
+    except OSError as exc:
+        assert exc.errno == errno.EBADF
+        return False
+
+
 class Stats:
     """
     Collects mean and variance, running average.

{returnn-1.20240626.135241 → returnn-1.20240626.231046}/returnn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240626.135241
+Version: 1.20240626.231046
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20240626.135241 → returnn-1.20240626.231046}/tests/test_rf_base.py

@@ -397,3 +397,43 @@ def test_rf_range_over_dim():
         rf.range_over_dim(time_dim, dtype="float32").mark_as_output("range_float", shape=[time_dim])
 
     run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
+def test_build_dict():
+    d = rf.build_dict(rf.BatchNorm, use_mask=True)
+    assert d == {"class": "rf.BatchNorm", "use_mask": True}
+
+
+def test_build_dict_func():
+    d = rf.build_dict(rf.relu)
+    assert d == {"class": "rf.relu"}
+
+
+def test_build_from_dict():
+    rf.select_backend_torch()
+    in_dim = Dim(7, name="in")
+    mod = rf.build_from_dict({"class": "rf.BatchNorm", "use_mask": True}, in_dim, eps=1e-4)
+    assert isinstance(mod, rf.BatchNorm)
+    assert mod.in_dim == in_dim
+    assert mod.use_mask
+    assert mod.eps == 1e-4
+
+
+def test_build_from_dict_func():
+    import functools
+
+    func = rf.build_from_dict({"class": "rf.relu"})
+    assert func is rf.relu
+    func = rf.build_from_dict({"class": "rf.combine", "kind": "+", "b": 1})
+    assert isinstance(func, functools.partial)
+    assert func.func is rf.combine
+    assert func.keywords == {"kind": "+", "b": 1}
+
+
+def test_build_from_dict_func_native():
+    from types import BuiltinFunctionType
+
+    rf.select_backend_torch()  # enables some of the native optimizations
+    assert isinstance(rf.combine, BuiltinFunctionType)  # due to native optimizations
+    func = rf.build_from_dict({"class": "rf.combine"})
+    assert func is rf.combine

returnn-1.20240626.135241/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240626.135241'
-long_version = '1.20240626.135241+git.744699a'