returnn 1.20250109.114342__tar.gz → 1.20250110.101532__tar.gz

{returnn-1.20250109.114342/returnn.egg-info → returnn-1.20250110.101532}/PKG-INFO

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

returnn-1.20250110.101532/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250110.101532'
+long_version = '1.20250110.101532+git.826ed76'

{returnn-1.20250109.114342 → returnn-1.20250110.101532}/returnn/datasets/distrib_files.py

@@ -364,8 +364,7 @@ class DistributeFilesDataset(CachedDataset2):
         Distributes the files from files_order into ``num_bins`` while attempting
         to make every bin as evenly sized (based on ``file_sizes``) as possible.
         """
-
-        total_size = sum(file_sizes.values())
+        total_size = sum(file_sizes[_get_key_for_file_tree(f_tree)] for f_tree in files_order)
         avg_size_per_sub_epoch = total_size / num_bins
         # Now evenly distribute the files over the bins.
         # Note that many one-pass variants of algorithms to evenly distribute

{returnn-1.20250109.114342 → returnn-1.20250110.101532}/returnn/frontend/decoder/transformer.py

@@ -203,6 +203,7 @@ class TransformerDecoder(rf.Module):
         state: rf.State,
         encoder: Optional[rf.State] = None,
         collected_outputs: Optional[Dict[str, Tensor]] = None,
+        output_only_last_frame: bool = False,
     ) -> Tuple[Tensor, rf.State]:
         """
         forward, single step or whole sequence.
@@ -212,6 +213,8 @@ class TransformerDecoder(rf.Module):
         :param state: e.g. via :func:`default_initial_state`
         :param encoder: via :func:`transform_encoder`
         :param collected_outputs:
+        :param output_only_last_frame: if True, and spatial_dim is not single_step_dim,
+            the returned logits will only be for the last frame
         :return: logits, new state
         """
         new_state = rf.State()
@@ -236,6 +239,9 @@ class TransformerDecoder(rf.Module):
             if collected_outputs is not None:
                 collected_outputs[layer_name] = decoded
 
+        if output_only_last_frame and spatial_dim != single_step_dim:
+            decoded = rf.gather(decoded, axis=spatial_dim, indices=rf.last_frame_position_of_dim(spatial_dim))
+
         decoded = self.final_layer_norm(decoded)
         logits = self.logits(decoded)
 

{returnn-1.20250109.114342 → returnn-1.20250110.101532}/returnn/frontend/dims.py

@@ -19,6 +19,7 @@ __all__ = [
     "dim_match_priority_when_needed",
     "num_elements_of_shape",
     "masked_fraction_of_shape",
+    "last_frame_position_of_dim",
 ]
 
 
@@ -195,3 +196,28 @@ def masked_fraction_of_shape(dims: Union[Dim, Sequence[Dim]], *, inverse: bool =
     for dim in dims:
         num_elems_total *= dim.get_dim_value_tensor()
     return (num_elems_masked / num_elems_total) if not inverse else (num_elems_total / num_elems_masked)
+
+
+def last_frame_position_of_dim(
+    dim: Dim, *, device: Optional[str] = None, allow_scalar_on_cpu: bool = True
+) -> Union[int, Tensor]:
+    """
+    :param dim:
+    :param device:
+    :param allow_scalar_on_cpu: if device is not given, we would use rf.get_default_device() for dynamic sizes.
+        If this is True, we would allow to return a scalar dynamic size on CPU.
+    :return: last frame position of dim
+    """
+    if dim.size is not None:
+        if dim.size > 0:
+            return dim.size - 1
+        return 0
+    if device is None:
+        if allow_scalar_on_cpu and not dim.dyn_size_ext.dims:
+            device = "cpu"
+        else:
+            device = rf.get_default_device()
+    pos = dim.get_dyn_size_ext_for_device(device) - 1
+    pos = rf.maximum(pos, 0)
+    pos.sparse_dim = dim
+    return pos

{returnn-1.20250109.114342 → returnn-1.20250110.101532/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250109.114342 → returnn-1.20250110.101532}/tests/test_Dataset.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Iterator, List, Dict
+from typing import Any, Iterator, List, Dict, Optional
 import os
 import sys
 import _setup_test_env  # noqa
@@ -905,8 +905,12 @@ def test_MapDatasetWrapper():
 def test_DistributeFilesDataset_distribute_evenly_by_size():
     from returnn.datasets.distrib_files import DistributeFilesDataset
 
-    def _test(sizes: List[int], partition_epoch: int, expected: List[List[int]]):
-        files = [f"file-{i}" for i in range(len(sizes))]
+    def _test(
+        sizes: List[int], partition_epoch: int, expected: List[List[int]], files_order: Optional[List[str]] = None
+    ):
+        files = files_order
+        if files is None:
+            files = [f"file-{i}" for i in range(len(sizes))]
         file_sizes = {f: s for f, s in zip(files, sizes)}
         res = DistributeFilesDataset._distribute_evenly_by_size(
             num_bins=partition_epoch, file_sizes=file_sizes, files_order=files
@@ -926,6 +930,7 @@ def test_DistributeFilesDataset_distribute_evenly_by_size():
         8,
         [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]],
     )
+    _test([1], 5, [[1, 1]] * 5, files_order=["file"] * 10)  # test duplicate files
 
     def _test_stddev(sizes: List[int], partition_epoch: int, max_stddev_percent: float):
         avg_per_bin = sum(sizes) / partition_epoch

returnn-1.20250109.114342/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250109.114342'
-long_version = '1.20250109.114342+git.bbbe0ae'