returnn 1.20251013.113026__py3-none-any.whl → 1.20260109.93428__py3-none-any.whl

returnn/frontend/conv.py

@@ -3,7 +3,7 @@ Convolution, transposed convolution, pooling
 """
 
 from __future__ import annotations
-from typing import Optional, Sequence, Tuple, Union
+from typing import Optional, Union, TypeVar, Sequence, Tuple, List
 from returnn.util.basic import next_type_attrib_in_mro_chain
 from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
@@ -25,6 +25,9 @@ __all__ = [
     "pool2d",
     "pool3d",
     "make_conv_out_spatial_dims",
+    "calc_conv_out_length",
+    "make_transposed_conv_out_spatial_dims",
+    "calc_transposed_conv_out_length",
 ]
 
 
@@ -396,7 +399,11 @@ def transposed_conv(
             )
         if use_mask:
             source = source.copy_masked(0, dims=in_spatial_dims)
-    if padding == "same" and _any_is_non_default(strides, default=1) and _should_use_consistent_same_padding():
+    if (
+        padding == "same"
+        and any(s != 1 for s in (strides or [fs.dimension for fs in filter_size]))
+        and _should_use_consistent_same_padding()
+    ):
         # I don't really know what this should mean here... Investigate this further...
         raise NotImplementedError("consistent same padding not implemented for transposed conv")
     # noinspection PyProtectedMember
@@ -424,6 +431,39 @@ class TransposedConv1d(_TransposedConv):
 
     nd = 1
 
+    def __init__(
+        self,
+        in_dim: Dim,
+        out_dim: Dim,
+        filter_size: Union[int, Dim],
+        *,
+        padding: str,
+        remove_padding: int = 0,
+        output_padding: Optional[int] = None,
+        strides: Optional[int] = None,
+        with_bias: bool = True,
+    ):
+        """
+        :param in_dim:
+        :param out_dim:
+        :param filter_size:
+        :param strides: specifies the upscaling. by default, same as filter_size
+        :param padding: "same" or "valid"
+        :param remove_padding:
+        :param output_padding:
+        :param with_bias: whether to add a bias. enabled by default
+        """
+        super().__init__(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            filter_size=[filter_size],
+            padding=padding,
+            remove_padding=remove_padding,
+            output_padding=output_padding,
+            strides=[strides] if strides is not None else None,
+            with_bias=with_bias,
+        )
+
     __call__ = _ConvOrTransposedConv._call_nd1
 
 
@@ -704,7 +744,7 @@ def make_conv_out_spatial_dims(
     strides: Union[Sequence[int], int] = 1,
     dilation_rate: Union[Sequence[int], int] = 1,
     description_prefix: Optional[str] = None,
-) -> Sequence[Dim]:
+) -> List[Dim]:
     """create out spatial dims from in spatial dims"""
     nd = len(in_spatial_dims)
     if isinstance(filter_size, (int, Dim)):
@@ -715,84 +755,263 @@ def make_conv_out_spatial_dims(
         strides = [strides] * nd
     if isinstance(dilation_rate, int):
         dilation_rate = [dilation_rate] * nd
-    assert nd == len(in_spatial_dims) == len(filter_size) == len(strides) == len(dilation_rate)
     if isinstance(padding, (int, str)):
         padding = [padding] * nd
+    assert nd == len(in_spatial_dims) == len(filter_size) == len(strides) == len(dilation_rate) == len(padding)
     padding = [p.lower() if isinstance(p, str) else p for p in padding]
     out_spatial_dims = []
     for i in range(nd):
-        in_spatial_dim = in_spatial_dims[i]
-        if (filter_size[i] == strides[i] == 1 and padding[i] in ("valid", "same", 0)) or (
-            strides[i] == 1 and padding[i] == "same"
-        ):
-            out_spatial_dims.append(in_spatial_dim)
-        else:
-            out_spatial_dim = _calc_out_dim(
-                in_dim=in_spatial_dim,
+        out_spatial_dims.append(
+            calc_conv_out_length(
+                in_spatial_dims[i],
                 filter_size=filter_size[i],
+                padding=padding[i],
                 stride=strides[i],
                 dilation_rate=dilation_rate[i],
-                padding=padding[i],
+                name=f"{description_prefix}:spatial{i}" if description_prefix else None,
             )
-            assert isinstance(out_spatial_dim, Dim)
-            if description_prefix and out_spatial_dim != in_spatial_dim:
-                out_spatial_dim.name = f"{description_prefix}:spatial{i}"
-            if in_spatial_dim.dyn_size_ext is not None and out_spatial_dim.dyn_size_ext is None:
-                out_spatial_dim.dyn_size_ext = _calc_out_dim(
-                    in_dim=in_spatial_dim.dyn_size_ext,
-                    filter_size=filter_size[i],
-                    stride=strides[i],
-                    dilation_rate=dilation_rate[i],
-                    padding=padding[i],
-                )
-            out_spatial_dims.append(out_spatial_dim)
+        )
     return out_spatial_dims
 
 
-def _calc_out_dim(in_dim, filter_size, stride, padding, dilation_rate=1):
+T = TypeVar("T", int, Dim, Tensor)
+
+
+def calc_conv_out_length(
+    in_length: Union[T, int, Dim, Tensor],
+    *,
+    filter_size: Union[T, int, Dim, Tensor],
+    stride: int,
+    padding: Union[str, int],
+    dilation_rate: int = 1,
+    name: Optional[str] = None,
+) -> T:
     """
     Copied and adapted from TF ConvLayer.calc_out_dim.
 
-    :param T|int|Tensor|torch.Tensor|tensorflow.Tensor|Dim in_dim: dimension in some axis
-    :param int filter_size: e.g. 2, for the corresponding axis
-    :param int stride: e.g. 1, for the corresponding axis
-    :param int dilation_rate: e.g. 1
-    :param str|int padding: "valid" or "same" or int
+    :param T in_length: dimension in some axis
+    :param filter_size: e.g. 2, for the corresponding axis
+    :param stride: e.g. 1, for the corresponding axis
+    :param dilation_rate: e.g. 1
+    :param padding: "valid" or "same" or int
+    :param name:
     :return: the output dimension
-    :rtype: T
     """
+    padding = padding.lower() if isinstance(padding, str) else padding
+    if isinstance(filter_size, int):
+        filter_size_int = filter_size
+    elif isinstance(filter_size, Dim):
+        filter_size_int = filter_size.dimension
+    else:
+        filter_size_int = None
+    filter_size_ = filter_size_int if isinstance(filter_size_int, int) else filter_size
 
-    def ceildiv(a, b):
-        """
-        :param T|int|Tensor|torch.Tensor|tensorflow.Tensor a:
-        :param T|int|Tensor|torch.Tensor|tensorflow.Tensor b:
-        :rtype: T
-        """
-        if isinstance(b, int) and b == 1:
-            return a
-        if isinstance(a, Tensor):
-            return rf.ceil_divide(a, b)
-        return -(-a // b)
+    if (filter_size_int == stride == 1 and padding in ("valid", "same", 0)) or (stride == 1 and padding == "same"):
+        return in_length
 
-    padding = padding.lower() if isinstance(padding, str) else padding
     # See tf.compat.v1.nn.convolution() documentation for more.
     if padding == "same":
-        if isinstance(in_dim, Dim):
-            return in_dim.ceildiv_right(stride)
-        return ceildiv(in_dim, stride)
+        if isinstance(in_length, Dim):
+            out_length = in_length.ceildiv_right(stride)
+        else:
+            out_length = _ceildiv(in_length, stride)
     elif padding == "valid" or isinstance(padding, int):
         if isinstance(padding, int) and padding != 0:
             assert padding > 0
-            in_dim = padding + in_dim + padding
-        if isinstance(in_dim, Dim):
-            filter_left_dilated = (filter_size - 1) * dilation_rate // 2
-            filter_right_dilated = (filter_size - 1) * dilation_rate - filter_left_dilated
-            valid_part = in_dim.sub_left(filter_left_dilated).sub_right(filter_right_dilated)
-            return valid_part.ceildiv_right(stride)
-        return ceildiv(in_dim - (filter_size - 1) * dilation_rate, stride)
+            in_length = padding + in_length + padding
+
+        if filter_size_int == 1:
+            valid_part = in_length
+        elif isinstance(in_length, Dim):
+            filter_left_dilated = (filter_size_ - 1) * dilation_rate // 2
+            filter_right_dilated = (filter_size_ - 1) * dilation_rate - filter_left_dilated
+            valid_part = in_length.sub_left(filter_left_dilated).sub_right(filter_right_dilated)
+        else:
+            valid_part = in_length - (filter_size_ - 1) * dilation_rate
+
+        if isinstance(valid_part, Dim):
+            out_length = valid_part.ceildiv_right(stride)
+        else:
+            out_length = _ceildiv(valid_part, stride)
+
     else:
         raise ValueError(f"invalid padding {padding!r} (type {type(padding).__name__})")
 
+    if isinstance(in_length, Dim):
+        assert isinstance(out_length, Dim)
+        if name and out_length != in_length:
+            out_length.name = name
+        if in_length.dyn_size_ext is not None and out_length.dyn_size_ext is None:
+            out_dyn_size_ext = calc_conv_out_length(
+                in_length=in_length.dyn_size_ext,
+                filter_size=filter_size,
+                stride=stride,
+                dilation_rate=dilation_rate,
+                padding=padding,
+            )
+            assert isinstance(out_dyn_size_ext, Tensor)
+            out_length.dyn_size_ext = out_dyn_size_ext
+
+    return out_length
+
+
+def make_transposed_conv_out_spatial_dims(
+    in_spatial_dims: Sequence[Dim],
+    *,
+    filter_size: Union[Sequence[Union[int, Dim]], int, Dim],
+    padding: Union[str, int, Sequence[int]],
+    output_padding: Optional[Union[Sequence[Optional[int]], int]] = None,
+    strides: Union[Sequence[Optional[int]], None, int] = None,
+    dilation_rate: Union[Sequence[int], int] = 1,
+    description_prefix: Optional[str] = None,
+) -> List[Dim]:
+    """create out spatial dims from in spatial dims"""
+    nd = len(in_spatial_dims)
+    if isinstance(filter_size, (int, Dim)):
+        filter_size = [filter_size] * nd
+    filter_size = [d.dimension if isinstance(d, Dim) else d for d in filter_size]
+    assert all(isinstance(s, int) for s in filter_size)
+    if isinstance(strides, int) or strides is None:
+        strides = [strides] * nd
+    if isinstance(dilation_rate, int):
+        dilation_rate = [dilation_rate] * nd
+    if isinstance(padding, (int, str)):
+        padding = [padding] * nd
+    if isinstance(output_padding, int) or output_padding is None:
+        output_padding = [output_padding] * nd
+    assert (
+        nd
+        == len(in_spatial_dims)
+        == len(filter_size)
+        == len(strides)
+        == len(dilation_rate)
+        == len(padding)
+        == len(output_padding)
+    )
+    padding = [p.lower() if isinstance(p, str) else p for p in padding]
+    out_spatial_dims = []
+    for i in range(nd):
+        out_spatial_dims.append(
+            calc_transposed_conv_out_length(
+                in_spatial_dims[i],
+                filter_size=filter_size[i],
+                padding=padding[i],
+                stride=strides[i],
+                dilation_rate=dilation_rate[i],
+                name=f"{description_prefix}:spatial{i}" if description_prefix else None,
+            )
+        )
+    return out_spatial_dims
+
+
+def calc_transposed_conv_out_length(
+    in_length: Union[T, int, Dim, Tensor],
+    *,
+    filter_size: Union[int, Dim],
+    padding: Union[int, str],
+    output_padding: Optional[int] = None,
+    stride: Optional[int] = None,
+    dilation_rate: int = 1,
+    name: Optional[str] = None,
+) -> T:
+    """
+    Determines output length of a transposed convolution given input length.
+
+    Copied from TF/Keras conv_utils.deconv_output_length
+    (https://github.com/tensorflow/tensorflow/blob/5912f51d580551e5cee2cfde4cb882594b4d3e60/tensorflow/python/keras/utils/conv_utils.py#L140),
+    adapted with simplification.
+
+    Also see :func:`calc_conv_out_length`.
+
+    :param in_length:
+    :param filter_size:
+    :param padding: one of `"same"`, `"valid"`, `"full"`.
+    :param output_padding: amount of padding along the output dimension.
+        Can be set to `None` in which case the output length is inferred.
+    :param stride:
+    :param dilation_rate:
+    :param name:
+    :returns: The output length (integer)
+    """
+    assert padding in {"same", "valid", "full"} or isinstance(padding, int)
+
+    if isinstance(filter_size, int):
+        filter_size_int = filter_size
+    elif isinstance(filter_size, Dim):
+        filter_size_int = filter_size.dimension
+    else:
+        filter_size_int = None
+    filter_size_ = filter_size_int if isinstance(filter_size_int, int) else filter_size
+
+    # Get the dilated kernel size
+    if dilation_rate != 1 and filter_size_int != 1:
+        filter_size = filter_size + (filter_size_ - 1) * (dilation_rate - 1)
+
+    if stride is None:
+        assert filter_size_int is not None
+        stride = filter_size_int
+    if stride != 1:
+        in_length = in_length * stride
+
+    # Infer length if output padding is None, else compute the exact length
+    if output_padding is None:
+        if padding == "valid" or padding == 0:
+            if filter_size_int == stride:
+                out_length = in_length
+            elif filter_size_int is not None:
+                out_length = in_length + max(filter_size_int - stride, 0)
+            elif isinstance(filter_size, Tensor):
+                out_length = in_length + rf.relu(filter_size - stride)
+            elif isinstance(filter_size, Dim):
+                out_length = in_length + (filter_size - stride)
+            else:
+                raise ValueError(f"invalid filter_size {filter_size!r} type {type(filter_size)}")
+        elif padding == "full":
+            out_length = in_length - (stride + filter_size_ - 2)
+        elif padding == "same":
+            out_length = in_length
+        else:
+            raise ValueError(f"invalid padding {padding!r}")
+
+    else:  # output_padding
+        if padding == "same":
+            pad = filter_size // 2
+        elif padding == "valid":
+            pad = 0
+        elif padding == "full":
+            pad = filter_size - 1
+        elif isinstance(padding, int):
+            pad = padding
+        else:
+            raise ValueError(f"invalid padding {padding!r}")
+        out_length = in_length + (filter_size - stride - 2 * pad + output_padding)
+
+    if isinstance(in_length, Dim):
+        assert isinstance(out_length, Dim)
+        if name and out_length != in_length:
+            out_length.name = name
+        if in_length.dyn_size_ext is not None and out_length.dyn_size_ext is None:
+            out_dyn_size_ext = calc_transposed_conv_out_length(
+                in_length=in_length.dyn_size_ext,
+                filter_size=filter_size,
+                padding=padding,
+                output_padding=output_padding,
+                stride=stride,
+                dilation_rate=dilation_rate,
+            )
+            assert isinstance(out_dyn_size_ext, Tensor)
+            out_length.dyn_size_ext = out_dyn_size_ext
+
+    return out_length
+
+
+def _ceildiv(a: T, b: Union[T, int, Tensor]) -> T:
+    if isinstance(b, int) and b == 1:
+        return a
+    if isinstance(a, Tensor):
+        return rf.ceil_divide(a, b)
+    return -(-a // b)
+
 
 def _should_use_consistent_same_padding() -> bool:
     """

returnn/frontend/device.py

@@ -8,7 +8,13 @@ from contextlib import contextmanager
 from returnn.tensor import Tensor
 
 
-__all__ = ["copy_to_device", "get_default_device", "set_default_device", "set_default_device_ctx"]
+__all__ = [
+    "copy_to_device",
+    "get_default_device",
+    "set_default_device",
+    "set_default_device_ctx",
+    "get_default_dim_size_device",
+]
 
 
 _default_device: Optional[str] = None
@@ -61,3 +67,10 @@ def set_default_device_ctx(device: Optional[str]):
         yield
     finally:
         _default_device = old_device
+
+
+def get_default_dim_size_device() -> Optional[str]:
+    """
+    :return: default device, where to put new tensors for dim sizes (Dim.dyn_size_ext)
+    """
+    return "cpu"

returnn/frontend/encoder/conformer.py

@@ -167,6 +167,25 @@ class ConformerConvSubsample(ISeqDownsamplingEncoder):
         out, _ = rf.merge_dims(x, dims=[self._final_second_spatial_dim, in_dim])
         return out, in_spatial_dims[0]
 
+    def get_out_spatial_dim(self, in_spatial_dim: Dim) -> Dim:
+        """Get output spatial dimension given input spatial dimension."""
+        out_spatial_dim = in_spatial_dim
+        for i, conv_layer in enumerate(self.conv_layers):
+            (out_spatial_dim,) = rf.make_conv_out_spatial_dims(
+                [out_spatial_dim],
+                filter_size=conv_layer.filter_size[0],
+                strides=conv_layer.strides[0],
+                padding=conv_layer.padding,
+            )
+            if self.pool_sizes and i < len(self.pool_sizes):
+                (out_spatial_dim,) = rf.make_conv_out_spatial_dims(
+                    [out_spatial_dim],
+                    filter_size=self.pool_sizes[i][0],
+                    strides=self.pool_sizes[i][0],
+                    padding="same",
+                )
+        return out_spatial_dim
+
 
 class ConformerEncoderLayer(rf.Module):
     """
@@ -273,6 +292,7 @@ class ConformerEncoderLayer(rf.Module):
         x_mhsa = self.self_att(x_mhsa_ln, axis=spatial_dim)
         x_mhsa = rf.dropout(x_mhsa, self.dropout, axis=self.dropout_broadcast and self.out_dim)
         x_mhsa_out = x_mhsa + x_ffn1_out
+        del x_mhsa
 
         # Conv
         x_conv_ln = self.conv_layer_norm(x_mhsa_out)

returnn/frontend/encoder/transformer.py

@@ -79,6 +79,8 @@ class TransformerEncoder(rf.Module):
         self.model_dim = model_dim
         self.embed_dim = embed_dim
 
+        self.out_dim = self.model_dim  # alias. consistency, compatibility
+
         if input_embedding is None or isinstance(input_embedding, rf.Module):
             pass
         elif isinstance(input_embedding, type):

returnn/frontend/loss.py

@@ -3,11 +3,12 @@ Loss functions
 """
 
 from __future__ import annotations
+from typing import Optional, Tuple
 from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
 
 
-__all__ = ["cross_entropy", "ctc_loss", "edit_distance"]
+__all__ = ["cross_entropy", "ctc_loss", "ctc_greedy_decode", "edit_distance"]
 
 
 def cross_entropy(
@@ -93,6 +94,44 @@ def ctc_loss(
     )
 
 
+def ctc_greedy_decode(
+    logits: Tensor,
+    *,
+    in_spatial_dim: Dim,
+    blank_index: int,
+    out_spatial_dim: Optional[Dim] = None,
+    target_dim: Optional[Dim] = None,
+    wb_target_dim: Optional[Dim] = None,
+) -> Tuple[Tensor, Dim]:
+    """
+    Greedy CTC decode.
+
+    :return: (labels, out_spatial_dim)
+    """
+    if wb_target_dim is None:
+        assert logits.feature_dim
+        wb_target_dim = logits.feature_dim
+
+    labels = rf.reduce_argmax(logits, axis=wb_target_dim)
+    labels = rf.cast(labels, "int32")
+
+    labels_shifted = rf.shift_right(labels, axis=in_spatial_dim, pad_value=blank_index)
+    mask_repeat = labels != labels_shifted
+    labels, out_spatial_dim = rf.masked_select(
+        labels,
+        mask=(labels != blank_index) & mask_repeat,
+        dims=[in_spatial_dim],
+        out_dim=out_spatial_dim,
+    )
+
+    if target_dim:
+        # Set correct sparse_dim. Only currently implemented if blank comes after.
+        assert target_dim.dimension == blank_index
+        labels.sparse_dim = target_dim
+
+    return labels, out_spatial_dim
+
+
 def edit_distance(a: Tensor, a_spatial_dim: Dim, b: Tensor, b_spatial_dim: Dim, *, dtype: str = "int32") -> Tensor:
     """
     :param a: [B,Ta]

returnn/frontend/math_.py

@@ -3,7 +3,6 @@ Math ops
 """
 
 from __future__ import annotations
-import typing
 from typing import Optional, Sequence, Union, Tuple, overload
 import numpy
 from returnn.tensor import Tensor, Dim
@@ -77,7 +76,7 @@ __all__ = [
 ]
 
 
-@typing.overload
+@overload
 def compare(
     a: Tensor,
     kind: str,
@@ -86,7 +85,19 @@ def compare(
     allow_broadcast_all_sources: Optional[bool] = None,
     dim_order: Optional[Sequence[Dim]] = None,
 ) -> Tensor:
-    """compare with two tensors"""
+    """compare"""
+
+
+@overload
+def compare(
+    a: Union[Tensor, _RawTensorTypes],
+    kind: str,
+    b: Union[Tensor, _RawTensorTypes],
+    *,
+    allow_broadcast_all_sources: Optional[bool] = None,
+    dim_order: Optional[Sequence[Dim]] = None,
+) -> Tensor:
+    """compare"""
 
 
 _CompareMap = {
@@ -138,7 +149,7 @@ def compare_bc(
     return compare(a, kind, b, allow_broadcast_all_sources=True, dim_order=dim_order)
 
 
-@typing.overload
+@overload
 def combine(
     a: Tensor,
     kind: str,
@@ -147,7 +158,19 @@ def combine(
     allow_broadcast_all_sources: Optional[bool] = None,
     dim_order: Optional[Sequence[Dim]] = None,
 ) -> Tensor:
-    """combine with two tensors"""
+    """combine"""
+
+
+@overload
+def combine(
+    a: Union[Tensor, _RawTensorTypes],
+    kind: str,
+    b: Union[Tensor, _RawTensorTypes],
+    *,
+    allow_broadcast_all_sources: Optional[bool] = None,
+    dim_order: Optional[Sequence[Dim]] = None,
+) -> Union[Tensor, _RawTensorTypes]:
+    """combine"""
 
 
 _CombineMap = {
@@ -332,7 +355,12 @@ def logical_not(a: Tensor) -> Tensor:
 
 @overload
 def opt_logical_or(a: bool, b: bool) -> bool:
-    """logical or"""
+    """opt logical or"""
+
+
+@overload
+def opt_logical_or(a: Union[Tensor, bool], b: Union[Tensor, bool]) -> Union[Tensor, bool]:
+    """opt logical or"""
 
 
 def opt_logical_or(a: Union[Tensor, bool], b: Union[Tensor, bool]) -> Union[Tensor, bool]:
@@ -350,7 +378,12 @@ def opt_logical_or(a: Union[Tensor, bool], b: Union[Tensor, bool]) -> Union[Tens
 
 @overload
 def opt_logical_and(a: bool, b: bool) -> bool:
-    """logical and"""
+    """opt logical and"""
+
+
+@overload
+def opt_logical_and(a: Union[Tensor, bool], b: Union[Tensor, bool]) -> Union[Tensor, bool]:
+    """opt logical and"""
 
 
 def opt_logical_and(a: Union[Tensor, bool], b: Union[Tensor, bool]) -> Union[Tensor, bool]:
@@ -416,16 +449,23 @@ def minimum(a: Tensor, b: Union[Tensor, _RawTensorTypes], *other_tensors) -> Ten
 
 def clip_by_value(
     x: Tensor,
-    clip_value_min: Union[Tensor, _RawTensorTypes],
-    clip_value_max: Union[Tensor, _RawTensorTypes],
+    clip_value_min: Union[None, Tensor, _RawTensorTypes] = None,
+    clip_value_max: Union[None, Tensor, _RawTensorTypes] = None,
     *,
     allow_broadcast_all_sources: bool = False,
 ) -> Tensor:
     """clip by value"""
-    # noinspection PyProtectedMember
-    return x._raw_backend.clip_by_value(
-        x, clip_value_min, clip_value_max, allow_broadcast_all_sources=allow_broadcast_all_sources
-    )
+    if clip_value_min is not None and clip_value_max is not None:
+        # noinspection PyProtectedMember
+        return x._raw_backend.clip_by_value(
+            x, clip_value_min, clip_value_max, allow_broadcast_all_sources=allow_broadcast_all_sources
+        )
+    elif clip_value_min is not None and clip_value_max is None:
+        return maximum(x, clip_value_min)
+    elif clip_value_min is None and clip_value_max is not None:
+        return minimum(x, clip_value_max)
+    else:
+        return x
 
 
 def identity(x: Tensor) -> Tensor:
@@ -541,7 +581,7 @@ def floor(a: Tensor) -> Tensor:
 
 # noinspection PyShadowingBuiltins
 def round(a: Tensor) -> Tensor:
-    """round"""
+    """round. the result dtype is same as input dtype, still float"""
     # noinspection PyProtectedMember
     return a._raw_backend.activation(a, "round")
 

returnn/frontend/module.py

@@ -274,10 +274,17 @@ class Functional(Module):
     (This is often not necessary, but sometimes useful.)
     """
 
-    def __init__(self, func):
+    def __init__(self, func, *, attribs: Optional[Dict[str, Any]] = None):
+        """
+        :param func: callable. you might want to use functools.partial if you want to fix some arguments.
+        :param attribs: optional dict of attributes to set on this module. e.g. ``out_dim``.
+        """
         super().__init__()
         assert callable(func)
         self.func = func
+        if attribs:
+            for k, v in attribs.items():
+                setattr(self, k, v)
 
     def __repr__(self):
         return f"{self.__class__.__name__}({self.func.__qualname__})"