returnn 1.20250123.234142__tar.gz → 1.20250124.221747__tar.gz

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

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

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

{returnn-1.20250123.234142 → returnn-1.20250124.221747}/returnn/frontend/dims.py

@@ -15,6 +15,7 @@ __all__ = [
     "range_over_dim_strided",
     "range_over_merged_dims",
     "replace_dim",
+    "replace_dim_v2",
     "set_sparse_dim",
     "dim_match_priority_when_needed",
     "num_elements_of_shape",
@@ -82,13 +83,15 @@ def range_over_merged_dims(
 
 def replace_dim(source: Tensor, *, in_dim: Dim, out_dim: Optional[Dim] = None) -> Tuple[Tensor, Dim]:
     """
-    Also see: :func:`rf.merge_dims`, :func:`rf.split_dims`.
+    Also see: :func:`replace_dim_v2`, :func:`rf.merge_dims`, :func:`rf.split_dims`.
 
     :param source:
-    :param in_dim:
-    :param out_dim:
-    :return: source with in_dim replaced by out_dim, and new out_dim.
-        this does not work for the sparse_dim. see :func:`set_sparse_dim` for that case.
+    :param in_dim: should be in ``source.dims``, to be replaced.
+        If you want to replace the ``source.sparse_dim``, see :func:`set_sparse_dim`.
+    :param out_dim: If not given, will create a new dim with the same size as ``in_dim``.
+        Note: If the size of ``out_dim`` is different from ``in_dim``,
+        currently the dim tag is replaced and there is no error -- this is not checked.
+    :return: ``source`` with ``in_dim`` replaced by ``out_dim``, and ``out_dim``.
     """
     if not out_dim:
         out_dim = in_dim.copy(same_as_self=False, description="new-dim")
@@ -96,6 +99,56 @@ def replace_dim(source: Tensor, *, in_dim: Dim, out_dim: Optional[Dim] = None) -
     return source._raw_backend.replace_dim(source, in_dim=in_dim, out_dim=out_dim), out_dim
 
 
+def replace_dim_v2(
+    source: Tensor, *, in_dim: Dim, out_dim: Dim, allow_expand: bool = True, allow_shrink: bool = True
+) -> Tensor:
+    """
+    Extends :func:`replace_dim` by also allowing to expand or shrink the dim
+    (or rather, to not ignore this; when :func:`replace_dim` is used on a dim with different size,
+     it will ignore this and anyway accept the new dim tag (currently)).
+
+    :param source:
+    :param in_dim: should be in ``source.dims``, to be replaced.
+        If you want to replace the ``source.sparse_dim``, see :func:`set_sparse_dim`.
+    :param out_dim: should not be in ``source.dims``, to be replaced.
+        Note: In contrast to :func:`replace_dim`, you must provide this explicitly.
+    :param allow_expand: if True, allow to expand the dim, i.e. if ``out_dim.size > in_dim.size``.
+    :param allow_shrink: if True, allow to shrink the dim, i.e. if ``out_dim.size < in_dim.size``.
+    :return: ``source`` with ``in_dim`` replaced by ``out_dim``.
+    """
+    if not rf.is_executing_eagerly():
+        raise NotImplementedError  # just not implemented yet. we can do via :func:`cond`
+    if in_dim not in source.dims:
+        raise ValueError(f"replace_dim_v2: dim {in_dim} not in {source}")
+    if out_dim in source.dims:
+        raise ValueError(f"replace_dim_v2: dim {out_dim} already in {source}")
+    old_size = in_dim.get_dim_value()
+    new_size = out_dim.get_dim_value()
+    if new_size == old_size:
+        res, _ = rf.replace_dim(source, in_dim=in_dim, out_dim=out_dim)
+    elif new_size > old_size:
+        if not allow_expand:
+            raise ValueError(
+                f"replace_dim_v2: not allowed to expand: {old_size} -> {new_size},"
+                f" for {in_dim=} {out_dim=}, in {source=}"
+            )
+        res, _ = rf.pad(
+            source,
+            axes=[in_dim],
+            padding=[(0, out_dim.get_dim_value_tensor() - in_dim.get_dim_value_tensor())],
+            out_dims=[out_dim],
+            value=0,
+        )
+    else:
+        if not allow_shrink:
+            raise ValueError(
+                f"replace_dim_v2: not allowed to shrink: {old_size} -> {new_size},"
+                f" for {in_dim=} {out_dim=}, in {source=}"
+            )
+        res, _ = rf.slice(source, axis=in_dim, size=out_dim)
+    return res
+
+
 def set_sparse_dim(source: Tensor, sparse_dim: Dim) -> Tensor:
     """
     :param source:

{returnn-1.20250123.234142 → returnn-1.20250124.221747}/returnn/torch/frontend/_backend.py

@@ -1767,19 +1767,25 @@ class TorchBackend(Backend[torch.Tensor]):
         source_raw = source.copy_compatible_to_dims_raw(source_templ_dims)
 
         out_dims = tuple(dims) + tuple(remaining_dims)
+        out_shape = [d.get_dim_value() for d in out_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
+            backup = backup.copy_transpose(out_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)
+        if torch.__version__ < (2, 1):
+            # There is a bug in older PyTorch where masked_scatter_ does not work correctly with non-contiguous tensors.
+            # https://github.com/pytorch/pytorch/issues/99638
+            out_raw = out_raw.contiguous()
+            mask_raw = mask_raw.contiguous()
+            source_raw = source_raw.contiguous()
         out_raw.masked_scatter_(mask_raw, source_raw)
         return Tensor(
             "masked_scatter",

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

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

{returnn-1.20250123.234142 → returnn-1.20250124.221747}/tests/test_rf_array.py

@@ -200,6 +200,145 @@ def test_pad_packed_batched():
     np.testing.assert_array_equal(in_.raw_tensor, out.raw_tensor)
 
 
+def test_masked_select_masked_scatter_vs_where_rev_dims():
+    """
+    Compare rf.where vs rf.masked_select+rf.masked_scatter.
+    Some op (e.g. LM update) could be done more efficiently on just the packed data (rf.masked_select),
+    that is why rf.masked_select+rf.masked_scatter can be useful over just using rf.where.
+    (In general, when computing the new ``a`` is expensive.)
+    The test does not cover this part on the computation (we just feed in some random ``a``),
+    but it checks that the results are the same,
+    as we had some problems with that in the past.
+    """
+    # noinspection PyShadowingNames
+    batch_dim = Dim(2, name="batch")
+    beam_dim = Dim(3, name="beam")
+
+    extern_data = TensorDict(
+        {
+            "mask": Tensor("mask", [beam_dim, batch_dim], dtype="bool"),
+            # Note: The dim order is relevant for this test. It passes when it is [batch_dim, beam_dim]...
+            "a": Tensor("a", [beam_dim, batch_dim], dtype="int32"),
+            "b": Tensor("b", [beam_dim, batch_dim], dtype="int32"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        mask = extern_data["mask"]
+        a = extern_data["a"]
+        b = extern_data["b"]
+        a.mark_as_output("a", shape=[beam_dim, batch_dim])
+        b.mark_as_output("b", shape=[beam_dim, batch_dim])
+        mask.mark_as_output("mask", shape=[beam_dim, batch_dim])
+
+        # Code via rf.where.
+        res_where = rf.where(mask, a, b)
+        assert res_where.dims_set == {beam_dim, batch_dim}
+        res_where.mark_as_output("res_where", shape=[beam_dim, batch_dim])
+
+        # Code via rf.masked_select and rf.masked_scatter.
+        a_packed, packed_dim = rf.masked_select(a, mask=mask, dims=[batch_dim, beam_dim])
+        assert a_packed.dims_set == {packed_dim}
+        res_packed = rf.masked_scatter(a_packed, b, mask=mask, dims=[batch_dim, beam_dim], in_dim=packed_dim)
+        assert res_packed.dims_set == {batch_dim, beam_dim}
+        res_packed.mark_as_output("res_packed", shape=[beam_dim, batch_dim])
+
+    out_dict = run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+    res_where = out_dict["res_where"]
+    res_packed = out_dict["res_packed"]
+    assert res_where.raw_tensor.shape == res_packed.raw_tensor.shape
+    print("a:")
+    print(out_dict["a"].raw_tensor)
+    print("b:")
+    print(out_dict["b"].raw_tensor)
+    print("mask:")
+    print(out_dict["mask"].raw_tensor)
+    print("result with where:")
+    print(res_where.raw_tensor)
+    print("result with packing:")
+    print(res_packed.raw_tensor)
+    np.testing.assert_equal(res_where.raw_tensor, res_packed.raw_tensor)
+
+
+def test_masked_select_masked_scatter_vs_where_md_rev_dims():
+    """
+    Like :func:`test_masked_select_masked_scatter_vs_where_rev_dims`
+    but we add another spatial dim, which then needs some further handling.
+    """
+    # noinspection PyShadowingNames
+    batch_dim = Dim(2, name="batch")
+    beam_dim = Dim(3, name="beam")
+    hist_a_dim = Dim(Tensor("hist_a", [batch_dim, beam_dim], dtype="int32"))
+    hist_b_dim = Dim(Tensor("hist_b", [batch_dim, beam_dim], dtype="int32"))
+
+    extern_data = TensorDict(
+        {
+            "mask": Tensor("mask", [beam_dim, batch_dim], dtype="bool"),
+            "a": Tensor("a", [beam_dim, batch_dim, hist_a_dim], dtype="int32"),
+            "b": Tensor("b", [beam_dim, batch_dim, hist_b_dim], dtype="int32"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        mask = extern_data["mask"]
+        a = extern_data["a"]
+        b = extern_data["b"]
+        a.mark_as_output("a", shape=[beam_dim, batch_dim, hist_a_dim])
+        b.mark_as_output("b", shape=[beam_dim, batch_dim, hist_b_dim])
+        mask.mark_as_output("mask", shape=[beam_dim, batch_dim])
+        hist_a_size = hist_a_dim.get_size_tensor()
+        hist_b_size = hist_b_dim.get_size_tensor()
+
+        # Code via rf.where.
+        hist_comb_sizes = rf.where(mask, hist_a_size, hist_b_size)
+        hist_comb_dim = Dim(hist_comb_sizes, name="hist_comb")
+        a_ = rf.replace_dim_v2(a, in_dim=hist_a_dim, out_dim=hist_comb_dim)
+        b_ = rf.replace_dim_v2(b, in_dim=hist_b_dim, out_dim=hist_comb_dim)
+        res_where = rf.where(mask, a_, b_)
+        assert res_where.dims_set == {beam_dim, batch_dim, hist_comb_dim}
+        res_where.mark_as_output("res_where", shape=[beam_dim, batch_dim, hist_comb_dim])
+
+        # Code via rf.masked_select and rf.masked_scatter.
+        hist_a_packed_size, packed_dim = rf.masked_select(hist_a_size, mask=mask, dims=[batch_dim, beam_dim])
+        hist_a_packed_dim = Dim(hist_a_packed_size, name="hist_a_packed")
+        a_packed, _ = rf.masked_select(a, mask=mask, dims=[batch_dim, beam_dim], out_dim=packed_dim)
+        a_packed = rf.replace_dim_v2(a_packed, in_dim=hist_a_dim, out_dim=hist_a_packed_dim)
+        assert a_packed.dims_set == {packed_dim, hist_a_packed_dim}
+        hist_merged_size = rf.masked_scatter(
+            hist_a_packed_size, hist_b_size, mask=mask, dims=[batch_dim, beam_dim], in_dim=packed_dim
+        )
+        hist_merged_dim = Dim(hist_merged_size, name="hist_merged")
+        a_packed = rf.replace_dim_v2(a_packed, in_dim=hist_a_packed_dim, out_dim=hist_merged_dim)
+        b_packed = rf.replace_dim_v2(b, in_dim=hist_b_dim, out_dim=hist_merged_dim)
+        res_packed = rf.masked_scatter(a_packed, b_packed, mask=mask, dims=[batch_dim, beam_dim], in_dim=packed_dim)
+        assert res_packed.dims_set == {batch_dim, beam_dim, hist_merged_dim}
+        res_packed.mark_as_output("res_packed", shape=[beam_dim, batch_dim, hist_merged_dim])
+
+    out_dict = run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+    res_where = out_dict["res_where"]
+    res_packed = out_dict["res_packed"]
+    hist_where_dim = res_where.dims[-1]
+    hist_packed_dim = res_packed.dims[-1]
+    assert hist_where_dim.dyn_size_ext.dims_set == hist_packed_dim.dyn_size_ext.dims_set == {batch_dim, beam_dim}
+    hist_where_size_raw = hist_where_dim.dyn_size_ext.copy_compatible_to_dims_raw((batch_dim, beam_dim))
+    hist_packed_size_raw = hist_packed_dim.dyn_size_ext.copy_compatible_to_dims_raw((batch_dim, beam_dim))
+    assert (hist_where_size_raw == hist_packed_size_raw).all()
+    assert res_where.raw_tensor.shape == res_packed.raw_tensor.shape
+    print("a:")
+    print(out_dict["a"].raw_tensor)
+    print("b:")
+    print(out_dict["b"].raw_tensor)
+    print("mask:")
+    print(out_dict["mask"].raw_tensor)
+    print("result with where:")
+    print(res_where.raw_tensor)
+    print("result with packing:")
+    print(res_packed.raw_tensor)
+    np.testing.assert_equal(res_where.raw_tensor, res_packed.raw_tensor)
+
+
 def test_reshape():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

returnn-1.20250123.234142/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250123.234142'
-long_version = '1.20250123.234142+git.65c8b3c'