returnn 1.20250116.102559__tar.gz → 1.20250117.155146__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250116.102559
+Version: 1.20250117.155146
 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.155146/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250117.155146'
+long_version = '1.20250117.155146+git.a41252b'

{returnn-1.20250116.102559 → returnn-1.20250117.155146}/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.155146/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.155146}/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.155146}/returnn/frontend/array_.py

@@ -578,7 +578,7 @@ def masked_select(
         return tensor._raw_backend.masked_select(tensor, mask=mask, dims=dims, out_dim=out_dim)
     # Separate implementation for the case where we have a subset of the mask dims, specifically one single dim.
     # See https://github.com/rwth-i6/returnn/issues/1605 for discussion.
-    mask = mask.copy_masked(mask_value=False, dims=dims)
+    mask = mask.copy_masked(mask_value=False)
     if len(dims) > 1:
         # Flatten it, in the specified order.
         tensor, in_dim = rf.merge_dims(tensor, dims=dims)
@@ -615,8 +615,33 @@ def masked_scatter(
     :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, backup=backup, mask=mask, dims=dims, in_dim=in_dim)
+    mask_dims_set = set(mask.dims)
+    dims_set = set(dims)
+    if not dims_set.issubset(mask_dims_set):
+        raise ValueError(f"masked_scatter: dims {dims} not subset of mask dims {mask.dims}")
+    if not dims_set:
+        raise ValueError(f"masked_scatter: dims {dims} empty")
+    if dims_set == mask_dims_set:
+        # noinspection PyProtectedMember
+        return source._raw_backend.masked_scatter(source, backup=backup, mask=mask, dims=dims, in_dim=in_dim)
+    # Separate implementation for the case where we have a subset of the mask dims.
+    # Keep this consistent to masked_select above.
+    mask = mask.copy_masked(mask_value=False)
+    if len(dims) > 1:
+        # Flatten it, in the specified order.
+        mask_, dim_ = rf.merge_dims(mask, dims=dims)
+    else:
+        mask_ = mask
+        (dim_,) = dims
+    dim_: Dim
+    idxs = rf.cumsum(rf.cast(mask_, "int32"), spatial_dim=dim_)  # [dim_] -> idx in dim_/in_dim + 1
+    idxs = rf.split_dims(idxs, dims=dims, axis=dim_)  # [dims...]
+    idxs = rf.where(mask, idxs - 1, 0)  # [dim_] -> idx in in_dim
+    res = rf.gather(source, axis=in_dim, indices=idxs)
+    if backup is None:
+        backup = 0
+    res = rf.where(mask, res, backup)
+    return res
 
 
 def sequence_mask(dims: Union[Dim, Sequence[Dim]], *, device: Optional[str] = None) -> Tensor:
@@ -733,6 +758,8 @@ def scatter(
     with mode=="sum",
     or otherwise it will take the max/min.
 
+    ``scatter`` is the inverse of :func:`gather`.
+
     (TF segment_sum can be implemented via this.)
 
     :param source: [batch_dims..., indices_dim(s)..., feature_dims...]

{returnn-1.20250116.102559 → returnn-1.20250117.155146}/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.155146}/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.155146}/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.155146}/returnn/torch/frontend/_backend.py

@@ -302,7 +302,6 @@ class TorchBackend(Backend[torch.Tensor]):
         pad_value: Union[None, int, float] = None,
     ) -> Tensor:
         """split dims"""
-        assert not axis.need_masking()  # not implemented
         assert pad_to_multiples in (None, False)  # not implemented
         axis_ = source.get_axis_from_description(axis)
         out_dims = source.dims[:axis_] + tuple(dims) + source.dims[axis_ + 1 :]
@@ -1747,6 +1746,11 @@ class TorchBackend(Backend[torch.Tensor]):
     ) -> Tensor:
         """masked scatter"""
         assert mask.dtype == "bool"
+        # Note: If mask.dims != dims, then sum(mask_raw.flatten()) could have less elements than source_raw
+        # (not counting remaining_dims), and then the out_raw.masked_scatter_ below fails silently!
+        # That's why we assert this here.
+        # Currently in the RF code, we have a generic fallback implementation,
+        # very similar to masked_select.
         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]
@@ -1765,6 +1769,7 @@ class TorchBackend(Backend[torch.Tensor]):
             out_dims = backup.dims
             out_raw = backup.raw_tensor.clone()  # we operate inplace below
 
+        mask = mask.copy_masked(mask_value=False)
         mask_raw = mask.copy_compatible_to_dims_raw(out_dims)
         out_raw.masked_scatter_(mask_raw, source_raw)
         return Tensor(

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250116.102559
+Version: 1.20250117.155146
 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.155146}/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