returnn 1.20250124.221747__tar.gz → 1.20250125.618__tar.gz

{returnn-1.20250124.221747/returnn.egg-info → returnn-1.20250125.618}/PKG-INFO

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

returnn-1.20250125.618/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250125.000618'
+long_version = '1.20250125.000618+git.b01684e'

{returnn-1.20250124.221747 → returnn-1.20250125.618}/returnn/frontend/__init__.py

@@ -56,3 +56,4 @@ from .types import *
 from . import audio
 from . import hooks
 from . import init
+from . import nested

returnn-1.20250125.618/returnn/frontend/nested.py

@@ -0,0 +1,311 @@
+"""
+Some utility functions on nested structures.
+"""
+
+from __future__ import annotations
+from typing import TypeVar, Optional, Sequence, Tuple, Dict
+import functools
+import re
+import tree
+from returnn.tensor import Tensor, Dim
+import returnn.frontend as rf
+
+
+__all__ = ["gather_nested", "masked_select_nested", "masked_scatter_nested"]
+
+
+T = TypeVar("T")
+
+
+def gather_nested(s: T, *, indices: Tensor, dim_map: Optional[Dict[Dim, Dim]] = None) -> T:
+    """
+    This is like :func:`gather`, but for nested structures.
+
+    :param s: nested structure
+    :param indices: indices tensor. see :func:`gather`
+    :param dim_map: if given, this will be updated with the new dim map
+    :return: s with gathered tensors
+    """
+    assert indices.sparse_dim
+    if dim_map is None:
+        dim_map = {}
+    tree.map_structure(functools.partial(_gather_prepare_dims, indices=indices, dim_map=dim_map), s)
+    s = tree.map_structure(functools.partial(_gather, indices=indices, dim_map=dim_map), s)
+    return s
+
+
+def _gather_prepare_dims(s: T, *, indices: Tensor, dim_map: Dict[Dim, Dim]) -> T:
+    if isinstance(s, Dim):
+        if s.dimension is not None:  # static
+            return s
+        if s in dim_map:
+            return dim_map[s]
+        if indices.sparse_dim in s.dyn_size_ext.dims:
+            new_dyn_size = _gather(s.dyn_size_ext, indices=indices, dim_map=dim_map)
+            new_dim = Dim(new_dyn_size, name=_extend_dim_name(s.name))
+            dim_map[s] = new_dim
+            return new_dim
+        return s
+    # everything else ignored at this stage
+
+
+def _gather(s: T, *, indices: Tensor, dim_map: Optional[Dict[Dim, Dim]] = None) -> T:
+    if isinstance(s, Tensor):
+        if dim_map and any(d in dim_map for d in s.dims):
+            for d in s.dims:
+                if d in dim_map:
+                    s = rf.replace_dim_v2(s, in_dim=d, out_dim=dim_map[d])
+        if indices.sparse_dim in s.dims:
+            # really the default case, otherwise e.g. scalar or so, independent of beam
+            s = rf.gather(s, indices=indices)
+        return s
+    if isinstance(s, Dim):
+        if s.dimension is not None:  # static
+            return s
+        if dim_map and s in dim_map:
+            return dim_map[s]
+        assert indices.sparse_dim not in s.dyn_size_ext.dims  # not expected, should be in dim_map
+        return s
+    raise TypeError(f"_gather: unexpected type ({type(s)})")
+
+
+def masked_select_nested(
+    s: T,
+    *,
+    mask: Tensor,
+    mask_cpu: Optional[Tensor] = None,
+    dims: Sequence[Dim],
+    out_dim: Optional[Dim] = None,
+    dim_map: Optional[Dict[Dim, Dim]] = None,
+) -> Tuple[T, Dim, Dict[Dim, Dim]]:
+    """
+    This is like :func:`masked_select`, but for nested structures.
+
+    :param s: nested structure
+    :param mask: mask tensor. see :func:`masked_select`
+    :param mask_cpu: mask tensor for CPU. this is used e.g. for dyn dim sizes
+    :param dims: dims to mask. see :func:`masked_select`
+    :param out_dim: the packed out dim. see :func:`masked_select`. if not given, a new one will be created.
+    :param dim_map: if given, this will be updated with the new dim map
+    :return: s with masked dims, out_dim, and a newly created dim map
+    """
+    if out_dim is None:
+        out_dim = Dim(None, name="packed_new_label")  # Flat_Batch_InBeam
+    if dim_map is None:
+        dim_map = {}
+    tree.map_structure(
+        functools.partial(
+            _masked_select_prepare_dims,
+            mask=mask_cpu if mask_cpu is not None else mask,
+            dims=dims,
+            out_dim=out_dim,
+            dim_map=dim_map,
+        ),
+        s,
+    )
+    s = tree.map_structure(
+        functools.partial(
+            _masked_select,
+            mask=mask,
+            mask_cpu=mask_cpu,
+            dims=dims,
+            out_dim=out_dim,
+            dim_map=dim_map,
+        ),
+        s,
+    )
+    return s, out_dim, dim_map
+
+
+def _masked_select_prepare_dims(s, *, mask: Tensor, dims: Sequence[Dim], out_dim: Dim, dim_map: Dict[Dim, Dim]):
+    if isinstance(s, Dim):
+        if s.dimension is not None:  # static
+            return s
+        if not any(d in s.dyn_size_ext.dims for d in dims):
+            return s
+        if s in dim_map:
+            return dim_map[s]
+        new_dyn_size = _masked_select(s.dyn_size_ext, mask=mask, dims=dims, out_dim=out_dim, dim_map=dim_map)
+        new_dim = Dim(new_dyn_size, name=_extend_dim_name(s.name))
+        dim_map[s] = new_dim
+        return new_dim
+    # everything else ignored at this stage
+
+
+def _masked_select(
+    s: T, *, mask: Tensor, mask_cpu: Optional[Tensor] = None, dims: Sequence[Dim], out_dim: Dim, dim_map: Dict[Dim, Dim]
+) -> T:
+    if isinstance(s, Tensor):
+        if not any(d in s.dims for d in dims):
+            return s  # e.g. scalar or so, independent from dims
+        if s.device == "cpu" and mask_cpu is not None:
+            mask = mask_cpu
+        # For the masked_select, we need that all masked dims are present, so add them if not.
+        # (E.g., when we mask [batch,beam], but we only have [batch], we need to add the beam dim.)
+        if any(d not in s.dims for d in dims):
+            s = rf.expand_dims(s, dims=[d for d in dims if d not in s.dims])
+        # The packing itself (masked_select).
+        s, _ = rf.masked_select(s, mask=mask, dims=dims, out_dim=out_dim)
+        # In the resulting tensor, potentially replace dims.
+        # In addition to the dim replacement, we also might need to slice, as the size might be smaller.
+        if any(d in dim_map for d in s.dims):
+            for d in s.dims:
+                if d in dim_map:
+                    s, _ = rf.slice(s, axis=d, size=dim_map[d])
+        return s
+    if isinstance(s, Dim):
+        if s.dimension is not None:  # static
+            return s
+        if not any(d in s.dyn_size_ext.dims for d in dims):
+            return s
+        assert s in dim_map
+        return dim_map[s]
+    raise TypeError(f"_masked_select: unexpected type ({type(s)})")
+
+
+def masked_scatter_nested(
+    s: T,
+    backup: T,
+    *,
+    mask: Tensor,
+    mask_cpu: Tensor,
+    dims: Sequence[Dim],
+    in_dim: Dim,
+    masked_select_dim_map: Dict[Dim, Dim],
+    masked_scatter_dim_map: Optional[Dict[Dim, Dim]] = None,
+) -> T:
+    """
+    Reverse of :func:`masked_select_nested`.
+
+    :param s: nested structure, where dims are packed, i.e. (in_dim,...)
+    :param backup: nested structure, where we scatter into. tensors like (dims...,...)
+    :param mask: mask tensor. see :func:`masked_scatter`/:func:`masked_select`
+    :param mask_cpu: mask tensor for CPU. this is used e.g. for dyn dim sizes. see :func:`masked_scatter`
+    :param dims: dims to mask. see :func:`masked_scatter`/:func:`masked_select`
+    :param in_dim: the packed in dim. see :func:`masked_scatter`
+    :param masked_select_dim_map: the dim map from :func:`masked_select_nested`.
+        This describes how to map dims from s to backup.
+    :param masked_scatter_dim_map: for any new dims created by this function, this will be updated
+    :return: backup with s scattered in
+    """
+    reverse_dim_map = {v: k for k, v in masked_select_dim_map.items()}
+    if masked_scatter_dim_map is None:
+        masked_scatter_dim_map = {}
+
+    tree.map_structure(
+        functools.partial(
+            _masked_scatter_merge_dims,
+            mask=mask_cpu,
+            dims=dims,
+            in_dim=in_dim,
+            reverse_dim_map=reverse_dim_map,
+            merged_dim_map=masked_scatter_dim_map,
+        ),
+        s,
+        backup,
+    )
+    s = tree.map_structure(
+        functools.partial(
+            _masked_scatter,
+            mask=mask,
+            mask_cpu=mask_cpu,
+            dims=dims,
+            in_dim=in_dim,
+            reverse_dim_map=reverse_dim_map,
+            merged_dim_map=masked_scatter_dim_map,
+        ),
+        s,
+        backup,
+    )
+    return s
+
+
+def _masked_scatter_merge_dims(
+    s: T,
+    backup: T,
+    *,
+    mask: Tensor,
+    dims: Sequence[Dim],
+    in_dim: Dim,
+    reverse_dim_map: Dict[Dim, Dim],
+    merged_dim_map: Dict[Dim, Dim],
+) -> T:
+    if isinstance(s, Dim):
+        # This is slightly more complex than in the _masked_select case:
+        # We need to merge the s and backup depending on the mask.
+        if s in reverse_dim_map:
+            s = reverse_dim_map[s]
+        if s == backup:
+            return s
+        if s in merged_dim_map:
+            return merged_dim_map[s]
+        # Note: s/backup might even be static dims.
+        new_size = _masked_scatter(
+            s.get_size_tensor(),
+            backup.get_size_tensor(),
+            mask=mask,
+            dims=dims,
+            in_dim=in_dim,
+            reverse_dim_map=reverse_dim_map,
+            merged_dim_map=merged_dim_map,
+        )
+        assert new_size.dims_set == (
+            (s.get_size_tensor().dims_set | backup.get_size_tensor().dims_set) - {in_dim}
+        ) | set(dims)
+        new_dim = Dim(new_size, name=backup.name)
+        merged_dim_map[s] = new_dim
+        merged_dim_map[backup] = new_dim
+        return new_dim
+    # everything else ignored at this stage
+
+
+def _masked_scatter(
+    s: T,
+    backup: T,
+    *,
+    mask: Tensor,
+    mask_cpu: Optional[Tensor] = None,
+    dims: Sequence[Dim],
+    in_dim: Dim,
+    reverse_dim_map: Dict[Dim, Dim],
+    merged_dim_map: Dict[Dim, Dim],
+) -> T:
+    if isinstance(s, Tensor):
+        assert isinstance(backup, Tensor)
+        if s.device == "cpu" and mask_cpu is not None:
+            mask = mask_cpu
+        if in_dim not in s.dims:
+            s = rf.expand_dim(s, in_dim)
+        # Do the reverse of _masked_select above.
+        # First replace the dims back.
+        if any(d in reverse_dim_map for d in s.dims):
+            for d in s.dims:
+                if d in reverse_dim_map:
+                    s = rf.replace_dim_v2(s, in_dim=d, out_dim=reverse_dim_map[d], allow_shrink=False)
+        # We also might need to replace newly merged dims, both in s and backup.
+        for d in s.dims:
+            if d in merged_dim_map:
+                s = rf.replace_dim_v2(s, in_dim=d, out_dim=merged_dim_map[d])
+        for d in backup.dims:
+            if d in merged_dim_map:
+                backup = rf.replace_dim_v2(backup, in_dim=d, out_dim=merged_dim_map[d])
+        # The unpacking itself (reversing the masked_select, i.e. masked_scatter).
+        s = rf.masked_scatter(s, backup, mask=mask, dims=dims, in_dim=in_dim)
+        return s
+    if isinstance(s, Dim):
+        # This is slightly more complex than in the _masked_select case:
+        # We need to merge the s and backup depending on the mask.
+        if s in reverse_dim_map:
+            s = reverse_dim_map[s]
+        if s in merged_dim_map:
+            return merged_dim_map[s]
+        return s
+    raise TypeError(f"_masked_scatter: unexpected type ({type(s)})")
+
+
+def _extend_dim_name(name: str) -> str:
+    # check ends with _<num>
+    m = re.match(r"^(.*)_(\d+)$", name)
+    if m:
+        return f"{m.group(1)}_{int(m.group(2)) + 1}"
+    return name + "_1"

{returnn-1.20250124.221747 → returnn-1.20250125.618/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250124.221747 → returnn-1.20250125.618}/returnn.egg-info/SOURCES.txt

@@ -186,6 +186,7 @@ returnn/frontend/loss.py
 returnn/frontend/math_.py
 returnn/frontend/matmul.py
 returnn/frontend/module.py
+returnn/frontend/nested.py
 returnn/frontend/normalization.py
 returnn/frontend/parameter.py
 returnn/frontend/parametrizations.py

returnn-1.20250124.221747/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250124.221747'
-long_version = '1.20250124.221747+git.fcc8ac0'