returnn 1.20240613.164354__tar.gz → 1.20240617.171744__tar.gz

{returnn-1.20240613.164354/returnn.egg-info → returnn-1.20240617.171744}/PKG-INFO

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

returnn-1.20240617.171744/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240617.171744'
+long_version = '1.20240617.171744+git.ef21456'

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/datasets/distrib_files.py

@@ -301,7 +301,7 @@ class DistributeFilesDataset(CachedDataset2):
         if "seq_ordering" not in dataset_dict and "seq_order_control_dataset" not in dataset_dict:
             raise ValueError(
                 f"{self}: sub dataset should have explicit seq_ordering "
-                "(or seq_order_control_dataset for MetaDataset), got: {dataset_dict}"
+                f"(or seq_order_control_dataset for MetaDataset), got: {dataset_dict}"
             )
         self._lazy_init_file_cache_proc()
         dataset_dict, exit_hook = self._file_cache.handle_cached_files_in_config(dataset_dict)

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/frontend/__init__.py

@@ -40,6 +40,7 @@ from .math_ import *
 from .matmul import *
 from .normalization import *
 from .parameter import *
+from .piecewise_linear import *
 from .rand import *
 from .rec import *
 from .reduce import *

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/frontend/_backend.py

@@ -135,14 +135,14 @@ class Backend(Generic[T]):
         raise NotImplementedError
 
     @staticmethod
-    def get_shape_raw(raw_tensor: T) -> Union[T, Tuple[Union[int, T]]]:
+    def get_shape_raw(raw_tensor: T) -> Union[T, Tuple[Union[int, T], ...]]:
         """
         :return: shape of raw tensor
         """
         raise NotImplementedError
 
     @staticmethod
-    def get_shape_tuple_raw(raw_tensor: T) -> Tuple[Union[int, T]]:
+    def get_shape_tuple_raw(raw_tensor: T) -> Tuple[Union[int, T], ...]:
         """
         :return: shape of raw tensor. assumes that ndim is known.
             In eager frameworks, all dims are int.
@@ -150,7 +150,7 @@ class Backend(Generic[T]):
         raise NotImplementedError
 
     @staticmethod
-    def get_known_shape_raw(raw_tensor: T) -> Tuple[Optional[int]]:
+    def get_known_shape_raw(raw_tensor: T) -> Tuple[Optional[int], ...]:
         """
         :return: shape of raw tensor, int for static known, None otherwise. assumes that ndim is known.
             This will not create any ops.
@@ -159,7 +159,7 @@ class Backend(Generic[T]):
         raise NotImplementedError
 
     @staticmethod
-    def set_known_shape_raw(raw_tensor: T, shape: Tuple[Optional[int]]) -> None:
+    def set_known_shape_raw(raw_tensor: T, shape: Tuple[Optional[int], ...]) -> None:
         """
         Sets the known shape of the raw tensor.
         This is only supported in graph-based frameworks,
@@ -506,6 +506,15 @@ class Backend(Generic[T]):
         """
         raise NotImplementedError
 
+    @staticmethod
+    def stack(sources: Sequence[Tensor], *, out_dim: Dim) -> Tensor:
+        """
+        :param sources:
+        :param out_dim:
+        :return: stacked tensor
+        """
+        raise NotImplementedError
+
     # Restrict the possible activation function names,
     # to not get unexpected behavior,
     # or unwanted incompatibilities.
@@ -953,6 +962,22 @@ class Backend(Generic[T]):
         """where"""
         raise NotImplementedError
 
+    @staticmethod
+    def search_sorted(
+        sorted_seq: Tensor, values: Tensor, *, axis: Dim, side: str = "left", out_dtype: str = "int32"
+    ) -> Tensor:
+        """
+        :param sorted_seq: [SharedDims...,axis], sequence of numbers, sorted low to high in the given axis.
+        :param values: [SharedDims...,OtherDims...], sequence of numbers to search for in ``sorted_seq``.
+        :param axis:
+        :param side: "left" or "right"
+        :param out_dtype:
+        :return: [SharedDims...,OtherDims...] -> axis, indices in axis in ``sorted_seq`` such that
+            sorted_seq[i-1] < value <= sorted_seq[i] if side=="left",
+            sorted_seq[i-1] <= value < sorted_seq[i] if side=="right".
+        """
+        raise NotImplementedError
+
     @staticmethod
     def clip_by_value(
         x: Tensor,
@@ -964,6 +989,25 @@ class Backend(Generic[T]):
         """clip by value"""
         raise NotImplementedError
 
+    @staticmethod
+    def lerp(
+        start: Tensor, end: Tensor, weight: Union[float, Tensor], *, allow_broadcast_all_sources: bool = False
+    ) -> Tensor:
+        """
+        Linear interpolation between start and end.
+        (Some backends might provide an optimized version of this.)
+
+        :param start:
+        :param end:
+        :param weight: scalar or tensor
+        :param allow_broadcast_all_sources:
+        :return: start + weight * (end - start)
+        """
+        # Default implementation.
+        if not allow_broadcast_all_sources:
+            return start + weight * (end - start)
+        return rf.combine_bc(start, "+", rf.combine_bc(weight, "*", rf.combine_bc(end, "-", start)))
+
     @staticmethod
     def matmul(a: Tensor[T], b: Tensor[T], *, reduce: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> Tensor[T]:
         """

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/frontend/_numpy_backend.py

@@ -52,14 +52,14 @@ class NumpyBackend(Backend[numpy.ndarray]):
         return raw_tensor.ndim
 
     @staticmethod
-    def get_shape_raw(raw_tensor: numpy.ndarray) -> Tuple[int]:
+    def get_shape_raw(raw_tensor: numpy.ndarray) -> Tuple[int, ...]:
         """
         :return: shape of raw tensor
         """
         return raw_tensor.shape
 
     @staticmethod
-    def get_shape_tuple_raw(raw_tensor: numpy.ndarray) -> Tuple[int]:
+    def get_shape_tuple_raw(raw_tensor: numpy.ndarray) -> Tuple[int, ...]:
         """
         :return: shape of raw tensor. assumes that ndim is known.
             In eager frameworks, all dims are int.
@@ -67,7 +67,7 @@ class NumpyBackend(Backend[numpy.ndarray]):
         return raw_tensor.shape
 
     @staticmethod
-    def get_known_shape_raw(raw_tensor: numpy.ndarray) -> Tuple[int]:
+    def get_known_shape_raw(raw_tensor: numpy.ndarray) -> Tuple[int, ...]:
         """
         :return: shape of raw tensor, int for static known, None otherwise. assumes that ndim is known.
             This will not create any ops.

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/frontend/array_.py

@@ -29,6 +29,7 @@ __all__ = [
     "concat_features",
     "pad",
     "cum_concat_step",
+    "stack",
     "masked_select",
     "masked_scatter",
     "sequence_mask",
@@ -39,6 +40,7 @@ __all__ = [
     "shift_right",
     "reverse_sequence",
     "where",
+    "search_sorted",
     "sparse_to_dense",
     "one_hot",
 ]
@@ -352,6 +354,13 @@ def concat(
 ) -> Tuple[Tensor, Dim]:
     """
     Concatenates multiple sources in the specified dimension.
+
+    Also see :func:`stack`.
+
+    :param sources: list of (tensor, dim) pairs. dim is the axis to concatenate on.
+    :param allow_broadcast: if True, the sources can have different dims, and the result will be broadcasted.
+    :param out_dim: reuse existing dim for the resulting concatenated dim, if given
+    :return: concatenated tensor, out_dim
     """
     assert sources
     if not allow_broadcast:
@@ -490,6 +499,23 @@ def cum_concat_step(
     )
 
 
+def stack(sources: Sequence[Tensor], *, out_dim: Optional[Dim] = None) -> Tuple[Tensor, Dim]:
+    """
+    Stack the sources in a new dimension.
+    All sources must have the same shape.
+
+    :param sources:
+    :param out_dim: if given, use this as the new dim
+    :return: stacked tensor, out_dim
+    """
+    if not sources:
+        raise ValueError("no sources to stack")
+    if not out_dim:
+        out_dim = Dim(len(sources), name="stack")
+    # noinspection PyProtectedMember
+    return sources[0]._raw_backend.stack(sources, out_dim=out_dim), out_dim
+
+
 def masked_select(
     tensor: Tensor, *, mask: Tensor, dims: Sequence[Dim], out_dim: Optional[Dim] = None
 ) -> Tuple[Tensor, Dim]:
@@ -749,6 +775,23 @@ def where(
     return cond._raw_backend.where(cond, true_, false_, allow_broadcast_all_sources=allow_broadcast_all_sources)
 
 
+def search_sorted(
+    sorted_seq: Tensor, values: Tensor, *, axis: Dim, side: str = "left", out_dtype: str = "int32"
+) -> Tensor:
+    """
+    :param sorted_seq: [SharedDims...,axis], sequence of numbers, sorted low to high in the given axis.
+    :param values: [SharedDims...,OtherDims...], sequence of numbers to search for in ``sorted_seq``.
+    :param axis:
+    :param side: "left" or "right"
+    :param out_dtype:
+    :return: [SharedDims...,OtherDims...] -> axis, indices in axis in ``sorted_seq`` such that
+        sorted_seq[i-1] < value <= sorted_seq[i] if side=="left",
+        sorted_seq[i-1] <= value < sorted_seq[i] if side=="right".
+    """
+    # noinspection PyProtectedMember
+    return sorted_seq._raw_backend.search_sorted(sorted_seq, values, axis=axis, side=side, out_dtype=out_dtype)
+
+
 def sparse_to_dense(
     labels: Union[Tensor, rf.RawTensorTypes],
     *,

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/frontend/math_.py

@@ -66,6 +66,7 @@ __all__ = [
     "softmax",
     "log_softmax",
     "gating",
+    "lerp",
 ]
 
 
@@ -567,3 +568,20 @@ def gating(
 
     a, b = rf.split(x, axis=axis, out_dims=[out_dim, out_dim])
     return act_func(a) * gate_func(b), out_dim
+
+
+def lerp(
+    start: Tensor, end: Tensor, weight: Union[float, Tensor], *, allow_broadcast_all_sources: bool = False
+) -> Tensor:
+    """
+    Linear interpolation between start and end.
+    (Some backends might provide an optimized version of this.)
+
+    :param start:
+    :param end:
+    :param weight: scalar or tensor
+    :param allow_broadcast_all_sources:
+    :return: start + weight * (end - start)
+    """
+    # noinspection PyProtectedMember
+    return start._raw_backend.lerp(start, end, weight, allow_broadcast_all_sources=allow_broadcast_all_sources)

returnn-1.20240617.171744/returnn/frontend/piecewise_linear.py

@@ -0,0 +1,55 @@
+"""
+Piecewise linear function
+"""
+
+from __future__ import annotations
+from typing import Union, Dict
+import numpy as np
+from returnn.tensor import Tensor, Dim
+import returnn.frontend as rf
+
+
+__all__ = ["PiecewiseLinear"]
+
+
+class PiecewiseLinear(rf.Module):
+    """
+    Piecewise linear function.
+    """
+
+    def __init__(self, points: Dict[Union[int, float], Union[float, Tensor]]):
+        """
+        :param points: dict of key -> value pairs.
+        """
+        super().__init__()
+        if not points:
+            raise ValueError(f"{self}: points must not be empty")
+        self._points_sorted = sorted(points.items())
+        self.points_dim = Dim(len(self._points_sorted), name="pcw_schd_pieces")
+        # Note: Use rf.Parameter to work around deepcopy issue. https://github.com/rwth-i6/returnn/issues/1541
+        self._keys = rf.Parameter(
+            rf.convert_to_tensor(
+                np.array([k for k, _ in self._points_sorted], dtype=rf.get_default_float_dtype()),
+                dims=[self.points_dim],
+            ),
+            auxiliary=True,
+        )
+        self._values = rf.Parameter(
+            rf.stack([rf.convert_to_tensor(v) for _, v in self._points_sorted], out_dim=self.points_dim)[0],
+            auxiliary=True,
+        )
+
+    def __call__(self, x: Tensor) -> Tensor:
+        """
+        :param x: (x_dims...) -> value in keys
+        :return: y: (x_dims...,y_dims...) -> value in values
+        """
+        index = rf.search_sorted(self._keys, x, axis=self.points_dim)
+        index = rf.clip_by_value(index, 1, self.points_dim.dimension - 1)
+        x_start = rf.gather(self._keys, indices=index - 1)
+        x_end = rf.gather(self._keys, indices=index)
+        x_frac = (x - x_start) / (x_end - x_start)
+        x_frac = rf.clip_by_value(x_frac, 0.0, 1.0)
+        y_start = rf.gather(self._values, indices=index - 1)
+        y_end = rf.gather(self._values, indices=index)
+        return rf.lerp(y_start, y_end, x_frac)

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/tensor/_dim_extra.py

@@ -4,7 +4,7 @@ or just rarely used attribs, such that we can save memory for the common case.
 """
 
 from __future__ import annotations
-from typing import TYPE_CHECKING, Optional, Union, Tuple, Sequence, Dict, List, Callable
+from typing import TYPE_CHECKING, Optional, Union, Any, Tuple, Sequence, Dict, List, Set, Callable
 import operator
 
 from returnn.util.basic import Entity
@@ -1286,15 +1286,15 @@ class _DimMixin:
         self: Dim,
         other: Dim,
         *,
-        ignore_feature_dim=False,
-        allow_same_feature_dim=False,
-        allow_same_spatial_dim=None,
-        treat_feature_as_spatial=False,
-        broadcast_matches=False,
-        unknown_spatial_matches=False,
-        undefined_matches=False,
-        derived_matches=False,
-        allow_old_behavior=False,
+        ignore_feature_dim: bool = False,
+        allow_same_feature_dim: bool = False,
+        allow_same_spatial_dim: Optional[bool] = None,
+        treat_feature_as_spatial: bool = False,
+        broadcast_matches: bool = False,
+        unknown_spatial_matches: bool = False,
+        undefined_matches: bool = False,
+        derived_matches: bool = False,
+        allow_old_behavior: bool = False,
     ) -> bool:
         """
         Compares self to other for equality.
@@ -1307,16 +1307,16 @@ class _DimMixin:
         and might potentially change in the future.
           https://github.com/rwth-i6/returnn/issues/634
 
-        :param Dim other:
-        :param bool ignore_feature_dim:
-        :param bool allow_same_feature_dim:
-        :param bool|None allow_same_spatial_dim:
-        :param bool treat_feature_as_spatial:
-        :param bool broadcast_matches:
-        :param bool unknown_spatial_matches:
-        :param bool undefined_matches:
-        :param bool derived_matches:
-        :param bool allow_old_behavior: useful e.g. for find_matching_dim_map
+        :param other:
+        :param ignore_feature_dim:
+        :param allow_same_feature_dim:
+        :param allow_same_spatial_dim:
+        :param treat_feature_as_spatial:
+        :param broadcast_matches:
+        :param unknown_spatial_matches:
+        :param undefined_matches:
+        :param derived_matches:
+        :param allow_old_behavior: useful e.g. for find_matching_dim_map
         """
         if self is other:  # first some fast path check
             return True
@@ -1820,12 +1820,13 @@ class _DimMixin:
         self._make_extra().copy_same_as = other
 
     @classmethod
-    def get_existing_tag_from_collection(cls, other, tags, is_equal_opts=None):
+    def get_existing_tag_from_collection(
+        cls, other: Dim, tags: Union[Sequence[Dim], Set[Dim]], is_equal_opts: Optional[Dict[str, Any]] = None
+    ) -> Optional[Dim]:
         """
-        :param Dim other:
-        :param list[Dim]|tuple[Dim]|set[Dim] tags:
-        :param dict[str]|None is_equal_opts: passed to Dim.is_equal
-        :rtype: Dim|None
+        :param other:
+        :param tags:
+        :param is_equal_opts: passed to Dim.is_equal
         """
         if is_equal_opts is None:
             is_equal_opts = {}
@@ -1842,13 +1843,17 @@ class _DimMixin:
         return None
 
     @classmethod
-    def get_all_dimension_tags(cls, data_list, is_equal_opts=None, unique_separate_axes=True):
-        """
-        :param list[_t.Tensor] data_list:
-        :param dict[str]|None is_equal_opts: passed to Dim.is_equal
-        :param bool unique_separate_axes: e.g. data_list=[Data with shape (B,5,5,10)] results in 4 dim tags, not 3.
+    def get_all_dimension_tags(
+        cls,
+        data_list: List[_t.Tensor],
+        is_equal_opts: Optional[Dict[str, Any]] = None,
+        unique_separate_axes: bool = True,
+    ) -> Tuple[List[Dim], util.DictRefKeys[_t.Tensor, List[Dim]]]:
+        """
+        :param data_list:
+        :param is_equal_opts: passed to Dim.is_equal
+        :param unique_separate_axes: e.g. data_list=[Data with shape (B,5,5,10)] results in 4 dim tags, not 3.
         :return: list of dimension tags, dict for data -> list of dimension tags (for each axis)
-        :rtype: (list[Dim], util.DictRefKeys[_t.Tensor, list[Dim]])
         """
         tags = []
         data_axes_dict = util.DictRefKeys()  # type: util.DictRefKeys[_t.Tensor, List[Dim]]
@@ -2323,11 +2328,11 @@ class Op:
     Op on :class:`Dim` which results in a derived :class:`Dim`.
     """
 
-    def __init__(self, kind, inputs, attribs=None):
+    def __init__(self, kind: str, inputs: List[Dim], attribs: Optional[Dict[str, Any]] = None):
         """
-        :param str kind: "add", "sub", "mul", "ceildiv"
-        :param list[Dim] inputs:
-        :param dict[str]|None attribs:
+        :param kind: "add", "sub", "mul", "ceildiv"
+        :param inputs:
+        :param attribs:
         """
         self.kind = kind
         self.inputs = inputs

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/tf/frontend_layers/_backend.py

@@ -151,12 +151,12 @@ class ReturnnLayersBackend(Backend[Layer]):
         raise NotImplementedError
 
     @staticmethod
-    def get_shape_tuple_raw(raw_tensor: Layer) -> Tuple[Union[int, Layer]]:
+    def get_shape_tuple_raw(raw_tensor: Layer) -> Tuple[Union[int, Layer], ...]:
         """shape"""
         raise NotImplementedError
 
     @staticmethod
-    def get_known_shape_raw(raw_tensor: Layer) -> Tuple[Optional[int]]:
+    def get_known_shape_raw(raw_tensor: Layer) -> Tuple[Optional[int], ...]:
         """known shape"""
         return raw_tensor.tensor.batch_shape
 

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/tf/frontend_low_level/_backend.py

@@ -72,7 +72,7 @@ class TFBackend(Backend[tf.Tensor]):
             return tf.shape(raw_tensor)
 
     @staticmethod
-    def get_shape_tuple_raw(raw_tensor: tf.Tensor) -> Tuple[Union[int, tf.Tensor]]:
+    def get_shape_tuple_raw(raw_tensor: tf.Tensor) -> Tuple[Union[int, tf.Tensor], ...]:
         """
         :return: shape of raw tensor. assumes that ndim is known
         """
@@ -87,14 +87,14 @@ class TFBackend(Backend[tf.Tensor]):
             return tuple(shape)
 
     @staticmethod
-    def get_known_shape_raw(raw_tensor: tf.Tensor) -> Tuple[Optional[int]]:
+    def get_known_shape_raw(raw_tensor: tf.Tensor) -> Tuple[Optional[int], ...]:
         """
         :return: shape of raw tensor, int for static known, None otherwise. assumes that ndim is known.
         """
         return tuple(raw_tensor.shape.as_list())
 
     @staticmethod
-    def set_known_shape_raw(raw_tensor: tf.Tensor, shape: Tuple[Optional[int]]) -> None:
+    def set_known_shape_raw(raw_tensor: tf.Tensor, shape: Tuple[Optional[int], ...]) -> None:
         """
         wrap tf.Tensor.set_shape
         """

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn/torch/frontend/_backend.py

@@ -109,19 +109,19 @@ class TorchBackend(Backend[torch.Tensor]):
         return raw_tensor.dim()
 
     @staticmethod
-    def get_shape_raw(raw_tensor: torch.Tensor) -> Tuple[int]:
+    def get_shape_raw(raw_tensor: torch.Tensor) -> Tuple[int, ...]:
         """shape"""
         return tuple(raw_tensor.shape)
 
     @staticmethod
-    def get_shape_tuple_raw(raw_tensor: torch.Tensor) -> Tuple[int]:
+    def get_shape_tuple_raw(raw_tensor: torch.Tensor) -> Tuple[int, ...]:
         """
         :return: shape of raw tensor
         """
         return tuple(raw_tensor.shape)
 
     @staticmethod
-    def get_known_shape_raw(raw_tensor: torch.Tensor) -> Tuple[Optional[int]]:
+    def get_known_shape_raw(raw_tensor: torch.Tensor) -> Tuple[Optional[int], ...]:
         """
         :return: shape of raw tensor; here for PyTorch the full shape is always known
         """
@@ -501,6 +501,14 @@ class TorchBackend(Backend[torch.Tensor]):
         out.raw_tensor = torch.cat((prev_accum.raw_tensor, source_raw), dim=prev_accum.get_axis_from_description(axis))
         return out
 
+    @staticmethod
+    def stack(sources: Sequence[Tensor], *, out_dim: Dim) -> Tensor:
+        """stack"""
+        out_dims = (out_dim,) + sources[0].dims
+        out = Tensor("stack", dims=out_dims, dtype=sources[0].dtype, sparse_dim=sources[0].sparse_dim)
+        out.raw_tensor = torch.stack([s.copy_compatible_to_dims_raw(out_dims[1:]) for s in sources], dim=0)
+        return out
+
     @staticmethod
     def activation_raw(raw_tensor: torch.Tensor, func: str) -> torch.Tensor:
         """
@@ -1094,6 +1102,38 @@ class TorchBackend(Backend[torch.Tensor]):
         out.raw_tensor = torch.where(cond_bc_raw, true_bc_raw, false_bc_raw)
         return out
 
+    @staticmethod
+    def search_sorted(
+        sorted_seq: Tensor, values: Tensor, *, axis: Dim, side: str = "left", out_dtype: str = "int32"
+    ) -> Tensor:
+        """search sorted"""
+        if out_dtype == "int32":
+            out_int32 = True
+        elif out_dtype == "int64":
+            out_int32 = False
+        else:
+            raise NotImplementedError(f"search_sorted: out_dtype {out_dtype} not supported")
+        if axis not in sorted_seq.dims:
+            raise ValueError(f"search_sorted: axis {axis} not in sorted_seqs {sorted_seq}")
+        if axis.need_masking():
+            raise NotImplementedError(f"search_sorted: dynamic axis {axis} not supported")
+        sorted_seq_dims = [dim for dim in sorted_seq.dims if dim != axis] + [axis]
+        for dim in sorted_seq_dims[:-1]:
+            if dim not in values.dims:
+                raise ValueError(f"search_sorted: dim {dim} in sorted_seq {sorted_seq} but not in values {values}")
+        values_rem_dims = [dim for dim in values.dims if dim not in sorted_seq_dims[:-1]]
+        values_dims = sorted_seq_dims[:-1] + values_rem_dims
+        sorted_seq_raw: torch.Tensor = sorted_seq.copy_compatible_to_dims_raw(sorted_seq_dims)
+        values_raw: torch.Tensor = values.copy_compatible_to_dims_raw(values_dims)
+        if len(values_rem_dims) != 1:
+            values_raw = values_raw.reshape(values_raw.shape[: len(sorted_seq_dims[:-1])] + (-1,))
+        out = Tensor("search_sorted", dims=sorted_seq_dims[:-1] + values_rem_dims, dtype=out_dtype, sparse_dim=axis)
+        out_raw = torch.searchsorted(sorted_seq_raw, values_raw, side=side, out_int32=out_int32)
+        if len(values_rem_dims) != 1:
+            out_raw = out_raw.reshape([dim.get_dim_value() for dim in out.dims])
+        out.raw_tensor = out_raw
+        return out
+
     @staticmethod
     def clip_by_value(
         x: Tensor,
@@ -1119,6 +1159,22 @@ class TorchBackend(Backend[torch.Tensor]):
         out.raw_tensor = torch.clamp(x_bc_raw, min_bc_raw, max_bc_raw)
         return out
 
+    @staticmethod
+    def lerp(
+        start: Tensor, end: Tensor, weight: Union[float, Tensor], *, allow_broadcast_all_sources: bool = False
+    ) -> Tensor:
+        """lerp"""
+        weight = rf.convert_to_tensor(weight, _backend=TorchBackend, device=start.device)
+        out = Tensor.get_common_data(
+            [start, end, weight], allow_broadcast_all_sources=allow_broadcast_all_sources, name="lerp"
+        )
+        out.raw_tensor = torch.lerp(
+            start.copy_compatible_to_dims_raw(out.dims),
+            end.copy_compatible_to_dims_raw(out.dims),
+            weight.copy_compatible_to_dims_raw(out.dims),
+        )
+        return out
+
     @staticmethod
     def matmul(a: _TT, b: _TT, *, reduce: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> _TT:
         """

returnn-1.20240617.171744/returnn/util/math.py

@@ -0,0 +1,28 @@
+"""
+Some mathematical functions, in pure NumPy.
+"""
+
+from __future__ import annotations
+from typing import Union, Dict
+import numpy
+
+
+def next_power_of_two(n: int) -> int:
+    """next power of two, >= n"""
+    return 2 ** (int(n - 1).bit_length())
+
+
+class PiecewiseLinear:
+    """
+    Piecewise linear function.
+    """
+
+    def __init__(self, values: Dict[Union[int, float], Union[int, float]]):
+        self._sorted_items = sorted(values.items())
+        self._sorted_keys = numpy.array([x for x, _ in self._sorted_items])
+        self._sorted_values = numpy.array([y for _, y in self._sorted_items])
+
+    def __call__(self, x: Union[int, float]) -> Union[int, float]:
+        steps = self._sorted_keys
+        values = self._sorted_values
+        return numpy.interp(x, steps, values)

{returnn-1.20240613.164354 → returnn-1.20240617.171744/returnn.egg-info}/PKG-INFO

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

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/returnn.egg-info/SOURCES.txt

@@ -184,6 +184,7 @@ returnn/frontend/matmul.py
 returnn/frontend/module.py
 returnn/frontend/normalization.py
 returnn/frontend/parameter.py
+returnn/frontend/piecewise_linear.py
 returnn/frontend/rand.py
 returnn/frontend/rec.py
 returnn/frontend/reduce.py
@@ -368,6 +369,7 @@ tests/test_rf_label_smoothing.py
 tests/test_rf_loop.py
 tests/test_rf_math.py
 tests/test_rf_normalization.py
+tests/test_rf_piecewise_linear.py
 tests/test_rf_rec.py
 tests/test_rf_reduce.py
 tests/test_rf_signal.py

{returnn-1.20240613.164354 → returnn-1.20240617.171744}/tests/test_Util.py

@@ -96,6 +96,26 @@ def test_slice_pad_zeros():
     assert_equal(list(slice_pad_zeros(np.array([1, 2, 3, 4]), begin=2, end=6)), [3, 4, 0, 0])
 
 
+def test_math_PiecewiseLinear():
+    from returnn.util.math import PiecewiseLinear
+
+    eps = 1e-5
+    f = PiecewiseLinear({1: 2, 3: 4, 5: 1})
+    assert_equal(f(0), 2)
+    assert_equal(f(1 - eps), 2)
+    assert_equal(f(1), 2)
+    assert_almost_equal(f(1 + eps), 2, decimal=4)
+    assert_equal(f(2), 3)
+    assert_almost_equal(f(3 - eps), 4, decimal=4)
+    assert_equal(f(3), 4)
+    assert_almost_equal(f(3 + eps), 4, decimal=4)
+    assert_equal(f(4), 2.5)
+    assert_almost_equal(f(5 - eps), 1, decimal=4)
+    assert_equal(f(5), 1)
+    assert_equal(f(5 + eps), 1)
+    assert_equal(f(6), 1)
+
+
 def test_parse_orthography_into_symbols():
     assert_equal(list("hi"), parse_orthography_into_symbols("hi"))
     assert_equal(list(" hello "), parse_orthography_into_symbols(" hello "))