returnn 1.20241217.122516__tar.gz → 1.20241219.123408__tar.gz

{returnn-1.20241217.122516/returnn.egg-info → returnn-1.20241219.123408}/PKG-INFO

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

returnn-1.20241219.123408/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241219.123408'
+long_version = '1.20241219.123408+git.037d844'

{returnn-1.20241217.122516 → returnn-1.20241219.123408}/returnn/__main__.py

@@ -577,16 +577,16 @@ def execute_main_task():
                 assert data, "set forward_data"
             else:
                 data = init_dataset(config.opt_typed_value("forward_data"))
-            # engine.epoch is usually the epoch of the loaded checkpoint,
-            # or what EngineBase.get_epoch_model will return.
-            # You can have both load and load_epoch, where load points to the checkpoint,
-            # and load_epoch is some other epoch, which you will get here for the dataset.
-            data.init_seq_order(epoch=engine.epoch or 1)
             forward_callback = config.typed_value("forward_callback")
             assert forward_callback, "no forward_callback specified"
             if callable(forward_callback):
                 forward_callback = forward_callback()
             allow_skipping_seqs = config.bool("allow_skipping_seqs_in_forward", False)
+            # engine.epoch is usually the epoch of the loaded checkpoint,
+            # or what EngineBase.get_epoch_model will return.
+            # You can have both load and load_epoch, where load points to the checkpoint,
+            # and load_epoch is some other epoch, which you will get here for the dataset.
+            # data.init_seq_order will be called inside. (Maybe the dataset is first moved to a subproc.)
             engine.forward_with_callback(
                 dataset=data, callback=forward_callback, allow_skipping_seqs=allow_skipping_seqs
             )

{returnn-1.20241217.122516 → returnn-1.20241219.123408}/returnn/frontend/_backend.py

@@ -478,7 +478,7 @@ class Backend(Generic[T]):
         source: Tensor,
         *,
         axes: Sequence[Dim],
-        padding: Sequence[Tuple[Union[Dim, int], Union[Dim, int]]],
+        padding: Sequence[Tuple[Union[Dim, int, Tensor], Union[Dim, int, Tensor]]],
         out_dims: Sequence[Dim],
         handle_dynamic_dims: bool,
         mode: str = "constant",
@@ -1165,11 +1165,14 @@ class Backend(Generic[T]):
         raise NotImplementedError
 
     @staticmethod
-    def masked_scatter(source: Tensor, *, mask: Tensor, dims: Sequence[Dim], in_dim: Dim) -> Tensor:
+    def masked_scatter(
+        source: Tensor, backup: Optional[Tensor] = None, *, mask: Tensor, dims: Sequence[Dim], in_dim: Dim
+    ) -> Tensor:
         """
         The inverse of :func:`masked_select`.
 
         :param source: [in_dim, F...]
+        :param backup: [dims..., F...] (or subset of those dims). zero if not given.
         :param mask: [dims...] -> bool (e.g. [B,T])
         :param dims: the order of the dims defines the format. those dims should be exactly the dims of the mask.
         :param in_dim: the dim of the source which should be scattered into the mask.

{returnn-1.20241217.122516 → returnn-1.20241219.123408}/returnn/frontend/array_.py

@@ -412,7 +412,7 @@ def pad(
     source: Tensor,
     *,
     axes: Sequence[Dim],
-    padding: Sequence[Tuple[Union[Dim, int], Union[Dim, int]]],
+    padding: Sequence[Tuple[Union[Dim, int, Tensor], Union[Dim, int, Tensor]]],
     out_dims: Optional[Sequence[Dim]] = None,
     mode: str = "constant",
     value: Optional[Union[rf.RawTensorTypes, Tensor]] = None,
@@ -436,13 +436,6 @@ def pad(
     if handle_dynamic_dims is None:
         handle_dynamic_dims = _pad_handle_dynamic_dims_default(axes, padding, mode=mode)
     if not out_dims:
-        if handle_dynamic_dims:
-            for left, right in padding:
-                if isinstance(left, Dim):
-                    assert not left.need_masking(), f"padding {padding} does not support dynamic left padding"
-                if isinstance(right, Dim):
-                    assert not right.need_masking(), f"padding {padding} does not support dynamic right padding"
-                # Note that even dynamic middle dims is not exactly correct...
         out_dims = [left + middle + right for middle, (left, right) in zip(axes, padding)]
     # noinspection PyProtectedMember
     return (
@@ -480,10 +473,7 @@ def _pad_handle_dynamic_dims_default(
     global _pad_handle_dynamic_dims_shown_warning
     if not _pad_handle_dynamic_dims_shown_warning:
         for middle, (left, right) in zip(pad_axes, padding):
-            middle: Dim
-            if not middle.need_masking() and (isinstance(left, int) or not left.need_masking()):
-                continue
-            if mode != "circular" and isinstance(right, int) and right == 0:
+            if not _pad_need_dyn_dim_handling(middle, left, right, mode=mode):
                 continue
 
             logging.getLogger("returnn.frontend").warning(
@@ -498,6 +488,20 @@ def _pad_handle_dynamic_dims_default(
     return False
 
 
+def _pad_need_dyn_dim_handling(
+    middle: Dim, left: Union[Dim, int, Tensor], right: Union[Dim, int, Tensor], *, mode: str
+) -> bool:
+    if (
+        not middle.need_masking()
+        and (isinstance(left, int) or (isinstance(left, Dim) and not left.need_masking()))
+        or (isinstance(left, Tensor) and not left.dims)
+    ):
+        return False
+    if mode != "circular" and isinstance(right, int) and right == 0:
+        return False
+    return True
+
+
 def cum_concat_step(
     source: Tensor, *, prev_accum: Tensor, axis: Dim, out_spatial_dim: Optional[Dim] = None
 ) -> Tuple[Tensor, Dim]:
@@ -598,18 +602,21 @@ def masked_select(
     return res, out_dim
 
 
-def masked_scatter(source: Tensor, *, mask: Tensor, dims: Sequence[Dim], in_dim: Dim) -> Tensor:
+def masked_scatter(
+    source: Tensor, backup: Optional[Tensor] = None, *, mask: Tensor, dims: Sequence[Dim], in_dim: Dim
+) -> Tensor:
     """
     The inverse of :func:`masked_select`.
 
     :param source: [in_dim, F...]
+    :param backup: [dims..., F...] (or subset of those dims). zero if not given.
     :param mask: [dims...] -> bool (e.g. [B,T])
     :param dims: the order of the dims defines the format. those dims should be exactly the dims of the mask.
     :param in_dim: the dim of the source which should be scattered into the mask.
     :return: [dims..., F...]
     """
     # noinspection PyProtectedMember
-    return source._raw_backend.masked_scatter(source, mask=mask, dims=dims, in_dim=in_dim)
+    return source._raw_backend.masked_scatter(source, backup=backup, mask=mask, dims=dims, in_dim=in_dim)
 
 
 def sequence_mask(dims: Union[Dim, Sequence[Dim]], *, device: Optional[str] = None) -> Tensor:

{returnn-1.20241217.122516 → returnn-1.20241219.123408}/returnn/tf/frontend_layers/_backend.py

@@ -366,7 +366,7 @@ class ReturnnLayersBackend(Backend[Layer]):
         source: Tensor,
         *,
         axes: Sequence[Dim],
-        padding: Sequence[Tuple[Union[Dim, int], Union[Dim, int]]],
+        padding: Sequence[Tuple[Union[Dim, int, Tensor], Union[Dim, int, Tensor]]],
         out_dims: Sequence[Dim],
         handle_dynamic_dims: bool,
         mode: str = "constant",

{returnn-1.20241217.122516 → returnn-1.20241219.123408}/returnn/torch/engine.py

@@ -1169,15 +1169,20 @@ class Engine(EngineBase):
         self._maybe_reset_dev_memory_caches()
         self._reset_dev_memory_stats()
 
-        if dataset.supports_seq_order_sorting():
-            # We can sort it. Sort it in reverse to make sure that we have enough memory right at the beginning.
-            print("Dataset supports sorting, i.e. it will be sorted for optimal performance.", file=log.v3)
-            dataset.seq_ordering = "sorted_reverse"
-        else:
-            print(
-                "Dataset does not support sorting, i.e. it will not be sorted for optimal performance.",
-                file=log.v3,
-            )
+        if dataset_init_epoch:
+            if not self.config.bool("sort_dataset", True):
+                pass
+            elif dataset.seq_ordering == "sorted_reverse":
+                pass
+            elif dataset.supports_seq_order_sorting():
+                # We can sort it. Sort it in reverse to make sure that we have enough memory right at the beginning.
+                print("Dataset supports sorting, i.e. it will be sorted for optimal performance.", file=log.v3)
+                dataset.seq_ordering = "sorted_reverse"
+            else:
+                print(
+                    "Dataset does not support sorting, i.e. it will not be sorted for optimal performance.",
+                    file=log.v3,
+                )
 
         if allow_skipping_seqs:
             # Dangerous! If you enable this, you could lose sequences,

{returnn-1.20241217.122516 → returnn-1.20241219.123408}/returnn/torch/frontend/_backend.py

@@ -444,7 +444,7 @@ class TorchBackend(Backend[torch.Tensor]):
         source: Tensor,
         *,
         axes: Sequence[Dim],
-        padding: Sequence[Tuple[Union[Dim, int], Union[Dim, int]]],
+        padding: Sequence[Tuple[Union[Dim, int, Tensor], Union[Dim, int, Tensor]]],
         out_dims: Sequence[Dim],
         handle_dynamic_dims: bool,
         mode: str = "constant",
@@ -459,14 +459,31 @@ class TorchBackend(Backend[torch.Tensor]):
                 raw_pad += [0, 0]
                 continue
             remaining_dims.remove(dim)
-            pad_ = padding[axes.index(dim)]
-            raw_pad += [
-                pad_[0].get_dim_value() if isinstance(pad_[0], Dim) else pad_[0],
-                pad_[1].get_dim_value() if isinstance(pad_[1], Dim) else pad_[1],
-            ]
+            left, right = padding[axes.index(dim)]
+            if isinstance(left, Dim):
+                if handle_dynamic_dims:
+                    assert not left.need_masking(), f"pad: left {left} needs masking, not supported currently"
+                left_ = left.get_dim_value()
+            elif isinstance(left, Tensor):
+                if handle_dynamic_dims:
+                    assert not left.dims, f"pad: left {left} needs masking, not supported currently"
+                left_ = torch.max(left.raw_tensor)
+            elif isinstance(left, int):
+                left_ = left
+            else:
+                raise TypeError(f"pad: invalid left {left!r}")
+            if isinstance(right, Dim):
+                right_ = right.get_dim_value()
+            elif isinstance(right, Tensor):
+                right_ = torch.max(right.raw_tensor)
+            elif isinstance(right, int):
+                right_ = right
+            else:
+                raise TypeError(f"pad: invalid right {right!r}")
+            raw_pad += [left_, right_]
             if not remaining_dims:
                 break
-        # Use torch.nn.functional.pad if possible.
+        # Use torch.nn.functional.pad if possible. It's possible for scalar `value` (or None).
         if (isinstance(value, Tensor) and value.dims == ()) or (not isinstance(value, Tensor)):
             if isinstance(value, Tensor):
                 assert value.dims == ()
@@ -476,7 +493,7 @@ class TorchBackend(Backend[torch.Tensor]):
             )
             out.raw_tensor = torch.nn.functional.pad(source.raw_tensor, pad=raw_pad, mode=mode, value=value)
         else:  # Fallback to concat.
-            assert isinstance(value, Tensor)
+            assert isinstance(value, Tensor) and value.dims  # non-scalar
             assert all(dim in source.dims and dim not in axes for dim in value.dims)
             assert len(axes) == 1  # not implemented otherwise currently...
             ext_dim = Dim(1, name="ext")
@@ -493,12 +510,23 @@ class TorchBackend(Backend[torch.Tensor]):
                         f"pad: mode {mode} not implemented with dynamic dims and handle_dynamic_dims=True"
                     )
                 for out_dim, middle, (left, right) in zip(out_dims, axes, padding):
-                    if middle.need_masking() or (isinstance(left, Dim) and left.need_masking()):
-                        if isinstance(right, Dim) or right > 0:
+                    if (
+                        middle.need_masking()
+                        or (isinstance(left, Dim) and left.need_masking())
+                        or (isinstance(left, Tensor) and left.dims)
+                    ):
+                        if isinstance(left, Dim):
+                            left = left.get_size_tensor(device=out.device)
+                        assert isinstance(left, (int, Tensor))
+                        if (
+                            isinstance(right, Dim)
+                            or (isinstance(right, int) and right > 0)
+                            or isinstance(right, Tensor)
+                        ):
                             mask = rf.compare_bc(
                                 rf.range_over_dim(out_dim, device=out.device),
                                 "<",
-                                rf.copy_to_device((left + middle).dyn_size_ext, out.device),
+                                left + middle.get_size_tensor(device=out.device),
                             )
                             out.raw_tensor = torch.where(
                                 mask.copy_compatible_to_dims_raw(out.dims),
@@ -1714,24 +1742,37 @@ class TorchBackend(Backend[torch.Tensor]):
         return out, out_dim
 
     @staticmethod
-    def masked_scatter(source: Tensor, *, mask: Tensor, dims: Sequence[Dim], in_dim: Dim) -> Tensor:
+    def masked_scatter(
+        source: Tensor, backup: Optional[Tensor] = None, *, mask: Tensor, dims: Sequence[Dim], in_dim: Dim
+    ) -> Tensor:
         """masked scatter"""
         assert mask.dtype == "bool"
         assert set(mask.dims) == set(dims)
         assert in_dim in source.dims
         remaining_dims = [d for d in source.dims if d not in mask.dims and d != in_dim]
         source_templ_dims = (in_dim,) + tuple(remaining_dims)
-        tensor_templ_dims = tuple(dims) + tuple(remaining_dims)
         source_raw = source.copy_compatible_to_dims_raw(source_templ_dims)
-        mask_raw = mask.copy_compatible_to_dims_raw(tensor_templ_dims)
-        out_shape = [d.get_dim_value() for d in tensor_templ_dims]
-        out_raw = torch.zeros(out_shape, dtype=source_raw.dtype, device=source_raw.device)
+
+        out_dims = tuple(dims) + tuple(remaining_dims)
+        if backup is None:
+            out_shape = [d.get_dim_value() for d in out_dims]
+            out_raw = torch.zeros(out_shape, dtype=source_raw.dtype, device=source_raw.device)
+        else:
+            assert set(backup.dims).issubset(out_dims), f"backup dims {backup.dims} not subset of out dims {out_dims}"
+            for d in out_dims:
+                if d not in backup.dims:
+                    backup = rf.expand_dim(backup, dim=d)
+            out_dims = backup.dims
+            out_raw = backup.raw_tensor.clone()  # we operate inplace below
+
+        mask_raw = mask.copy_compatible_to_dims_raw(out_dims)
         out_raw.masked_scatter_(mask_raw, source_raw)
         return Tensor(
             "masked_scatter",
-            dims=tensor_templ_dims,
-            dtype=source.dtype,
+            dims=out_dims,
+            dtype=TorchBackend.get_dtype_name_raw(out_raw),
             sparse_dim=source.sparse_dim,
+            feature_dim=source.feature_dim,
             raw_tensor=out_raw,
         )
 

{returnn-1.20241217.122516 → returnn-1.20241219.123408/returnn.egg-info}/PKG-INFO

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

returnn-1.20241217.122516/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241217.122516'
-long_version = '1.20241217.122516+git.d51e906'