returnn 1.20250116.102559__tar.gz → 1.20250117.110003__tar.gz

{returnn-1.20250116.102559/returnn.egg-info → returnn-1.20250117.110003}/PKG-INFO

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

returnn-1.20250117.110003/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250117.110003'
+long_version = '1.20250117.110003+git.44fde7a'

{returnn-1.20250116.102559 → returnn-1.20250117.110003}/returnn/datasets/basic.py

@@ -1456,6 +1456,7 @@ def get_dataset_class(name: Union[str, Type[Dataset]]) -> Optional[Type[Dataset]
         "multi_proc",
         "distrib_files",
         "postprocessing",
+        "text_dict",
     ]
     for mod_name in mod_names:
         mod = import_module("returnn.datasets.%s" % mod_name)

returnn-1.20250117.110003/returnn/datasets/text_dict.py

@@ -0,0 +1,259 @@
+"""
+:class:`TextDictDataset`
+"""
+
+from __future__ import annotations
+from typing import Optional, Union, Any, Sequence, Tuple, List, Dict
+import numpy as np
+
+from returnn.log import log
+from .basic import DatasetSeq
+from .cached2 import CachedDataset2
+from .util.vocabulary import Vocabulary
+
+
+class TextDictDataset(CachedDataset2):
+    """
+    This dataset can read files in the format as usually generated from RETURNN search,
+    i.e. with beam like (item_format = "list_with_scores")::
+
+        {
+            seq_tag: [(score1, text1), (score2, text2), ...],
+            ...
+        }
+
+    Or without beam like (item_format = "single")::
+
+        {
+            seq_tag: text,
+            ...
+        }
+
+    The data keys:
+
+        data: The single (or best) sequence (encoded via vocab).
+        data_flat: for list_with_scores, all sequences concatenated (encoded via vocab), in the given order
+        data_seq_lens: for list_with_scores, the sequence lengths of each seq in data_flat
+        scores: for list_with_scores, the scores of each seq in data_flat
+    """
+
+    def __init__(
+        self,
+        *,
+        filename: str,
+        item_format: str = "list_with_scores",
+        vocab: Union[Vocabulary, Dict[str, Any]],
+        **kwargs,
+    ):
+        """
+        :param filename: text dict file. can be gzipped.
+        :param item_format: "list_with_scores" or "single"
+        :param vocab: to encode the text as a label sequence. See :class:`Vocabulary.create_vocab`.
+        """
+        super().__init__(**kwargs)
+        self.filename = filename
+        self.item_format = item_format
+        self.vocab = vocab if isinstance(vocab, Vocabulary) else Vocabulary.create_vocab(**vocab)
+        self.num_inputs = self.vocab.num_labels
+        self.num_outputs = {}
+        self.labels = {}
+        if item_format == "list_with_scores":
+            self.num_outputs.update(
+                {
+                    "data": (self.vocab.num_labels, 1),
+                    "data_flat": (self.vocab.num_labels, 1),
+                    "data_seq_lens": (1, 1),
+                    "scores": (1, 1),
+                }
+            )
+            self.labels.update({"data_flat": self.vocab.labels})
+        elif item_format == "single":
+            self.num_outputs.update({"data": (self.vocab.num_labels, 1)})
+            self.labels.update({"data": self.vocab.labels})
+        else:
+            raise ValueError(f"invalid item_format {item_format!r}")
+        self._data_values: Optional[List[Union[List[Tuple[float, str]], str]]] = None  # lazily loaded
+        self._seq_tags: Optional[List[str]] = None  # lazily loaded
+        self._seq_order: Optional[Sequence[int]] = None  # via init_seq_order
+
+    def _load(self):
+        if self._data_values is not None:
+            return
+
+        if self.filename.endswith(".gz"):
+            import gzip
+
+            txt = gzip.GzipFile(self.filename, "rb").read()
+        else:
+            txt = open(self.filename, "rb").read()
+
+        from returnn.util.literal_py_to_pickle import literal_eval
+
+        # Note: literal_py_to_pickle.literal_eval is quite efficient.
+        # However, currently, it does not support inf/nan literals,
+        # so it might break for some input.
+        # We might want to put a simple fallback to eval here if needed.
+        # Or maybe extend literal_py_to_pickle.literal_eval to support inf/nan literals.
+        try:
+            data: Dict[str, Any] = literal_eval(txt)
+        except Exception as exc:
+            print(f"{self}: Warning: literal_py_to_pickle.literal_eval failed:", file=log.v3)
+            print(f"  {type(exc).__name__}: {exc}", file=log.v3)
+            print("  Fallback to eval...", file=log.v3)
+            data: Dict[str, Any] = eval(txt)
+        assert data is not None
+        assert isinstance(data, dict)
+        assert len(data) > 0
+        # Check some data.
+        key, value = next(iter(data.items()))
+        assert isinstance(key, str), f"{self}: expected seq tag as keys, got {key!r} ({type(key)})"  # seq tag
+        if self.item_format == "single":
+            assert isinstance(value, str), f"{self}: expected str ({self.item_format}), got {value!r} ({type(value)})"
+        elif self.item_format == "list_with_scores":
+            assert isinstance(value, list), f"{self}: expected list ({self.item_format}), got {value!r} ({type(value)})"
+            assert len(value) > 0, f"{self}: expected non-empty list ({self.item_format}), got {value!r} for seq {key}"
+            value0 = value[0]
+            assert (
+                isinstance(value0, tuple)
+                and len(value0) == 2
+                and isinstance(value0[0], float)
+                and isinstance(value0[1], str)
+            ), f"{self}: expected (score,text) tuples ({self.item_format}), got {value0!r} ({type(value0)})"
+        else:
+            raise ValueError(f"invalid item_format {self.item_format!r}")
+        self._data_values = list(data.values())
+        self._seq_tags = list(data.keys())
+
+    def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
+        """init seq order"""
+        super().init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
+        if epoch is None and seq_list is None and seq_order is None:
+            self._num_seqs = 0
+            return True
+
+        random_seed = self._get_random_seed_for_epoch(epoch=epoch)
+        self.vocab.set_random_seed(random_seed)
+
+        if self.item_format == "single":
+
+            def _get_seq_len(i: int) -> int:
+                return len(self._data_values[i])
+
+        elif self.item_format == "list_with_scores":
+
+            def _get_seq_len(i: int) -> int:
+                values = self._data_values[i]
+                return sum(len(text) for _, text in values)
+
+        else:
+            raise ValueError(f"invalid item_format {self.item_format!r}")
+
+        if seq_order is not None:
+            self._seq_order = seq_order
+        elif seq_list is not None:
+            raise NotImplementedError(f"{self}: seq_list not supported yet")
+        else:
+            self._load()
+            num_seqs = len(self._data_values)
+            self._seq_order = self.get_seq_order_for_epoch(epoch=epoch, num_seqs=num_seqs, get_seq_len=_get_seq_len)
+        self._num_seqs = len(self._seq_order)
+
+    def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
+        orig_seq_idx = self._seq_order[seq_idx]
+        seq_tag = self._seq_tags[orig_seq_idx]
+        data = self._data_values[orig_seq_idx]
+
+        res = {}
+        labels_dtype = self.get_data_dtype("data")
+        if self.item_format == "single":
+            res["data"] = np.array(self.vocab.get_seq(data), dtype=labels_dtype)
+        elif self.item_format == "list_with_scores":
+            _, best = max(data, key=lambda x: x[0])
+            res["data"] = np.array(self.vocab.get_seq(best), dtype=labels_dtype)
+            labels = [self.vocab.get_seq(txt) for _, txt in data]
+            res["data_flat"] = np.array(sum(labels, []), dtype=labels_dtype)
+            res["data_seq_lens"] = np.array([len(seq) for seq in labels], dtype=self.get_data_dtype("data_seq_lens"))
+            res["scores"] = np.array([score for score, _ in data], dtype=self.get_data_dtype("scores"))
+        else:
+            raise ValueError(f"invalid item_format {self.item_format!r}")
+        return DatasetSeq(seq_idx=seq_idx, features=res, seq_tag=seq_tag)
+
+    def supports_sharding(self) -> bool:
+        """:return: whether this dataset supports sharding"""
+        return True
+
+    def supports_seq_order_sorting(self) -> bool:
+        """supports sorting"""
+        return True
+
+    def get_current_seq_order(self) -> Sequence[int]:
+        """:return: seq order"""
+        assert self._seq_order is not None, "init_seq_order not called"
+        return self._seq_order
+
+    def have_corpus_seq_idx(self) -> bool:
+        """
+        :return: whether we can use :func:`get_corpus_seq_idx`
+        """
+        return True
+
+    def get_corpus_seq_idx(self, seq_idx: int) -> int:
+        """
+        :param seq_idx:
+        """
+        assert self._seq_order is not None, "init_seq_order not called"
+        return self._seq_order[seq_idx]
+
+    def get_tag(self, seq_idx: int) -> str:
+        """
+        :param seq_idx:
+        :return: seq tag
+        """
+        self._load()
+        return self._seq_tags[self._seq_order[seq_idx]]
+
+    def get_all_tags(self) -> List[str]:
+        """:return: all tags"""
+        self._load()
+        return self._seq_tags
+
+    def get_total_num_seqs(self, *, fast: bool = False) -> int:
+        """:return: total num seqs in dataset (not for (sub)epoch)"""
+        self._load()
+        return len(self._data_values)
+
+    def get_data_dim(self, key: str) -> int:
+        """:return: dim of data entry with `key`"""
+        if key == "data" or key == "data_flat":
+            return self.vocab.num_labels
+        elif key == "data_seq_lens":
+            return 1
+        elif key == "scores":
+            return 1
+        else:
+            raise ValueError(f"{self}: unknown data key: {key}")
+
+    def get_data_dtype(self, key: str) -> str:
+        """:return: dtype of data entry with `key`"""
+        if key == "data" or key == "data_flat":
+            return "int32"
+        elif key == "data_seq_lens":
+            return "int32"
+        elif key == "scores":
+            return "float32"
+        else:
+            raise ValueError(f"{self}: unknown data key: {key}")
+
+    def get_data_keys(self) -> List[str]:
+        """:return: available data keys"""
+        return list(self.num_outputs.keys())
+
+    def get_data_shape(self, key: str) -> List[str]:
+        """
+        :returns get_data(*, key).shape[1:], i.e. num-frames excluded
+        """
+        return []  # all are scalar or sparse
+
+    def is_data_sparse(self, key: str) -> bool:
+        """:return: whether data entry with `key` is sparse"""
+        return key == "data" or key == "data_flat"

{returnn-1.20250116.102559 → returnn-1.20250117.110003}/returnn/datasets/util/vocabulary.py

@@ -123,8 +123,6 @@ class Vocabulary:
         Sets self.vocab, self.labels, self.num_labels.
         """
         filename = self.vocab_file
-        import pickle
-
         if self._labels is not None:
             self._vocab = {label: i for i, label in enumerate(self._labels)}
             self.num_labels = len(self._labels)
@@ -132,10 +130,17 @@ class Vocabulary:
             self._vocab, self._labels = self._cache[filename]
             self.num_labels = len(self._labels)
         else:
-            if filename[-4:] == ".pkl":
+            if filename.endswith(".pkl"):
+                import pickle
+
                 d = pickle.load(open(filename, "rb"))
             else:
-                file_content = open(filename, "r").read()
+                if filename.endswith(".gz"):
+                    import gzip
+
+                    file_content = gzip.open(filename, "rt").read()
+                else:
+                    file_content = open(filename, "r").read()
                 if file_content.startswith("{"):
                     d = eval(file_content)
                 else:

{returnn-1.20250116.102559 → returnn-1.20250117.110003}/returnn/tensor/_tensor_extra.py

@@ -174,23 +174,17 @@ class _TensorMixin(_TensorMixinBase):
             This is deprecated. Rather, the placeholder should be created outside and passed in.
         :param str|dict[str]|returnn.datasets.util.vocabulary.Vocabulary|None vocab: vocab of the feature dim
             or sparse dim.
-            This is deprecated. Rather, the vocab is part of the :class:`Dim`.
         :param dict[int|str,Dim]|None same_dim_tags_as: will mark our dimension tags to be the same
         """
         assert isinstance(self, _t.Tensor)
         shape, sparse, dim, batch_dim_axis, dim_tags  # noqa  # unused here, handled in infer_dim_tags
 
         if vocab is not None:
-            from returnn.datasets.util.vocabulary import Vocabulary
-
-            if isinstance(vocab, str):
-                vocab = Vocabulary(vocab)
-            elif isinstance(vocab, dict):
-                vocab = Vocabulary.create_vocab(**vocab)
-            assert isinstance(vocab, Vocabulary)
             assert self.sparse, "%s should represent indices of %s" % (self, vocab)
-            assert self.dim == vocab.num_labels, "%s dims do not match with vocab %s" % (self, vocab)
-            self.sparse_dim.vocab = vocab
+            if not self.sparse_dim.vocab:  # might already have been set earlier
+                vocab = _get_vocab(vocab)
+                assert self.dim == vocab.num_labels, "%s dims do not match with vocab %s" % (self, vocab)
+                self.sparse_dim.vocab = vocab
 
         if kwargs:
             self._extra = _TensorExtra(tensor=self, **kwargs)
@@ -3421,39 +3415,42 @@ def infer_sparse_dim(
     *,
     name: str,
     sparse: Optional[bool] = None,
-    sparse_dim,
     dim=NotSpecified,
+    vocab=None,
     **_other_kwargs,
 ) -> Optional[Dim]:
     """
+    Called when sparse_dim is None,
+    but we assume it is sparse
+
     :param name:
     :param sparse:
-    :param sparse_dim:
     :param dim:
+    :param vocab:
     :return: sparse dim
     """
     if sparse is None:
-        sparse = sparse_dim not in (None, NotSpecified)
-    if sparse_dim in (None, NotSpecified):
-        if sparse:
-            assert dim is not NotSpecified, "need dim (num classes) if sparse"
-            assert dim is None or isinstance(dim, int)
-            sparse_dim = Dim(
-                kind=Dim.Types.Feature,
-                dimension=dim,
-                description="%s:sparse-dim" % name,
-                auto_generated=True,
-            )
+        if vocab is None:
+            return None
+        sparse = True
+    assert isinstance(sparse, bool)
+    if not sparse:
+        return None
+    vocab = _get_vocab(vocab) if vocab else None
+    if vocab:
+        if dim is NotSpecified or dim is None:
+            dim = vocab.num_labels
         else:
-            sparse_dim = None
-    if sparse_dim is not None:
-        assert isinstance(sparse_dim, Dim)
-        assert sparse_dim.can_be_used_as_dim()
-        assert sparse
-        if dim is not NotSpecified:
-            assert sparse_dim.dimension == dim
-    else:
-        assert not sparse
+            assert dim == vocab.num_labels
+    assert dim is not NotSpecified, "need dim (num classes) if sparse"
+    assert dim is None or isinstance(dim, int)
+    sparse_dim = Dim(
+        kind=Dim.Types.Feature,
+        dimension=dim,
+        description="%s:sparse-dim" % name,
+        auto_generated=True,
+        vocab=vocab,
+    )
     return sparse_dim
 
 
@@ -3535,6 +3532,17 @@ def infer_dim_tags(
     return dims
 
 
+def _get_vocab(vocab):
+    from returnn.datasets.util.vocabulary import Vocabulary
+
+    if isinstance(vocab, str):
+        vocab = Vocabulary(vocab)
+    elif isinstance(vocab, dict):
+        vocab = Vocabulary.create_vocab(**vocab)
+    assert isinstance(vocab, Vocabulary)
+    return vocab
+
+
 class _SizePlaceholderProxy:
     """
     This is a proxy object to emulate the original Tensor.size_placeholder behavior,

{returnn-1.20250116.102559 → returnn-1.20250117.110003}/returnn/tensor/tensor.py

@@ -84,8 +84,8 @@ class Tensor(_TensorMixin, _TensorOpOverloadsMixin, Generic[RawTensorType]):
             - v2: time_dim_axis, feature_dim_axis are None by default.
         :param kwargs: see :func:`_handle_extra_kwargs`, :func:`infer_dim_tags`
         """
-        if "sparse" in kwargs and sparse_dim is None:
-            sparse_dim = _tensor_extra.infer_sparse_dim(name=name, sparse_dim=sparse_dim, **kwargs)
+        if ("sparse" in kwargs or "vocab" in kwargs) and sparse_dim is None:
+            sparse_dim = _tensor_extra.infer_sparse_dim(name=name, **kwargs)
         if dims is not None:
             assert "shape" not in kwargs and "dim_tags" not in kwargs  # probably old code got this wrong
             if version is None:

{returnn-1.20250116.102559 → returnn-1.20250117.110003}/returnn/tf/layers/basic.py

@@ -8,10 +8,11 @@ from typing import Optional, Union, Sequence, List, Tuple, Dict
 import typing
 import tensorflow as tf
 import contextlib
+from returnn.tensor import Tensor, Dim
 import returnn.tf.compat as tf_compat
 import returnn.tf.util.basic as tf_util
 from returnn.util.basic import unicode, NotSpecified
-from returnn.tf.util.data import Data, SearchBeam, Dim, FeatureDim, SpatialDim
+from returnn.tf.util.data import Data, SearchBeam, FeatureDim, SpatialDim
 from returnn.tf.util.basic import OutputWithActivation, dimshuffle, swapaxes
 from returnn.log import log
 from .base import LayerBase, Loss, InternalLayer, SearchChoices
@@ -8445,7 +8446,7 @@ class TimeChunkingLayer(_ConcatInputLayer):
         x = x.copy_with_batch_dim_axis(1)
         self.input_data = x
         in_dim = x.dim_tags[0]
-        x_t = x.placeholder
+        x_t = x.placeholder  # [T,B,...]
         if in_dim.dyn_size is not None:
             index = tf.cast(tf_util.sequence_mask_time_major(in_dim.dyn_size), tf.float32)
         else:
@@ -8458,13 +8459,14 @@ class TimeChunkingLayer(_ConcatInputLayer):
         from returnn.tf.native_op import chunk
 
         out, oindex = chunk(x_t, index=index, chunk_step=chunk_step, chunk_size=chunk_size)
+        # out shape (chunk_size, n_batch * n_chunks, n_dim), oindex shape (chunk_size, n_batch * n_chunks)
         if ext_rem_shape:
-            out = tf.reshape(out, tf.concat([tf.shape(oindex), ext_rem_shape], axis=0))
+            out = tf.reshape(out, tf.concat([tf.shape(oindex), ext_rem_shape], axis=0))  # [C_size,B*C,...]
         self.output.placeholder = out
         out.set_shape(self.output.batch_shape)
         out_dim = self.output.dim_tags[0]
         if out_dim.dimension is None and out_dim.dyn_size is None:
-            out_dim.dyn_size = tf.reduce_sum(tf.cast(oindex, tf.int32), axis=0)
+            out_dim.dyn_size = tf.reduce_sum(tf.cast(oindex, tf.int32), axis=0)  # [B*C]
 
     @classmethod
     def get_out_data_from_opts(cls, name, sources, axis="T", out_dim=None, **kwargs):
@@ -8480,8 +8482,20 @@ class TimeChunkingLayer(_ConcatInputLayer):
         in_dim = data.dim_tags[axis]
         data = data.copy_move_axis(old_axis=axis, new_axis=0)  # (T,...)
         data = data.copy_with_batch_dim_axis(1)  # (T,B,...)
+        old_batch_dim = data.get_batch_dim_tag()
+        new_batch_dim = Dim(
+            kind=old_batch_dim.kind,
+            description=f"{name}:chunked_batch",
+            dimension=Tensor(f"{name}:chunked_batch", [], "int32"),  # unknown here...
+        )
+        data = data.copy_template_replace_dim_tag(1, new_batch_dim)
         if not out_dim:
-            out_dim = Dim(kind=in_dim.kind, description="%s:chunking" % name, auto_generated=True, dimension=None)
+            out_dim = Dim(
+                Tensor(f"{name}:chunk_size", [new_batch_dim], "int32"),
+                kind=in_dim.kind,
+                description=f"{name}:chunking",
+                auto_generated=True,
+            )
         data = data.copy_template_replace_dim_tag(axis=0, new_dim_tag=out_dim)  # (T',B',...)
         data.time_dim_axis = 0
         return data

{returnn-1.20250116.102559 → returnn-1.20250117.110003/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250116.102559 → returnn-1.20250117.110003}/returnn.egg-info/SOURCES.txt

@@ -112,6 +112,7 @@ returnn/datasets/postprocessing.py
 returnn/datasets/raw_wav.py
 returnn/datasets/sprint.py
 returnn/datasets/stereo.py
+returnn/datasets/text_dict.py
 returnn/datasets/util/__init__.py
 returnn/datasets/util/feature_extraction.py
 returnn/datasets/util/strings.py

{returnn-1.20250116.102559 → returnn-1.20250117.110003}/tests/test_TFNetworkLayer.py

@@ -3313,18 +3313,25 @@ def test_out_shape():
 
 
 def _check_MergeDimsLayer(
-    session, in_data_opts, in_static_shape, opts, out_data_shape, out_static_shape, in_sizes=None, out_sizes=None
-):
+    session: tf.compat.v1.Session,
+    in_data_opts: Dict[str, Any],
+    in_static_shape: Tuple[int, ...],
+    opts: Dict[str, Any],
+    out_data_shape: Tuple[Optional[int], ...],
+    out_static_shape: Tuple[int, ...],
+    in_sizes: Optional[Dict[int, Tuple[int, ...]]] = None,
+    out_sizes: Optional[Dict[int, Tuple[int, ...]]] = None,
+) -> MergeDimsLayer:
     """
-    :param tf.compat.v1.Session session:
-    :param dict[str] in_data_opts:
-    :param tuple[int] in_static_shape:
-    :param dict[str] opts: for MergeDimsLayer
-    :param tuple[int|None] out_data_shape:
-    :param tuple[int] out_static_shape:
-    :param dict[int,tuple[int]]|None in_sizes:
-    :param dict[int,tuple[int]]|None out_sizes:
-    :rtype: MergeDimsLayer
+    :param session:
+    :param in_data_opts:
+    :param in_static_shape:
+    :param opts: for MergeDimsLayer
+    :param out_data_shape:
+    :param out_static_shape:
+    :param in_sizes:
+    :param out_sizes:
+    :return: layer
     """
     net = TFNetwork(extern_data=ExternData())
     rnd = numpy.random.RandomState(42)
@@ -9479,6 +9486,63 @@ def test_PostfixInTimeLayer():
                 assert out[0, src_seq_lens[0] + repeat - 1, 0] == -7
 
 
+def test_TimeChunkingLayer():
+    n_batch, n_time, n_in = 2, 11, 3
+    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
+    in_seq_lens = numpy.array([11, 9])
+    config = Config({"extern_data": {"data": {"shape": (None, n_in)}}})
+    with make_scope() as session:
+        net = TFNetwork(config=config)
+        net.construct_from_dict(
+            {"output": {"class": "time_chunking", "chunk_size": 5, "chunk_step": 5, "from": "data"}}
+        )
+        out = net.get_default_output_layer().output
+        print("out:", out)
+        out_v, out_lens = session.run(
+            (out.placeholder, out.get_sequence_lengths()),
+            feed_dict={
+                net.extern_data.get_batch_info().dim: n_batch,
+                net.extern_data.data["data"].placeholder: in_v,
+                net.extern_data.data["data"].get_sequence_lengths(): in_seq_lens,
+            },
+        )
+        assert isinstance(out_v, numpy.ndarray)
+        print(out_v.shape)
+
+
+def test_TimeChunkingLayer_TimeUnchunkingLayer():
+    n_batch, n_time, n_in = 2, 11, 3
+    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
+    in_seq_lens = numpy.array([11, 9])
+    for b in range(n_batch):
+        in_v[b, in_seq_lens[b] :] = 0
+    config = Config({"extern_data": {"data": {"shape": (None, n_in)}}})
+    with make_scope() as session:
+        net = TFNetwork(config=config)
+        net.construct_from_dict(
+            {
+                "chunked": {"class": "time_chunking", "chunk_size": 5, "chunk_step": 5, "from": "data"},
+                "output": {"class": "time_unchunking", "chunking_layer": "chunked", "from": "chunked"},
+            }
+        )
+        in_ = net.get_layer("data").output
+        out = net.get_default_output_layer().output
+        print("out:", out)
+        out = out.copy_transpose(in_.dims).copy_masked(0.0)
+        out_v, out_lens = session.run(
+            (out.placeholder, out.get_sequence_lengths()),
+            feed_dict={
+                net.extern_data.get_batch_info().dim: n_batch,
+                net.extern_data.data["data"].placeholder: in_v,
+                net.extern_data.data["data"].get_sequence_lengths(): in_seq_lens,
+            },
+        )
+        assert isinstance(out_v, numpy.ndarray)
+        assert out_v.shape == in_v.shape
+        print(out_v)
+        numpy.testing.assert_equal(out_v, in_v)
+
+
 def test_DotLayer():
     with make_scope() as session:
         B = 2

returnn-1.20250116.102559/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250116.102559'
-long_version = '1.20250116.102559+git.dcae190'