returnn 1.20240824.1611__tar.gz → 1.20240826.132634__tar.gz

{returnn-1.20240824.1611 → returnn-1.20240826.132634}/PKG-INFO

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

returnn-1.20240826.132634/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240826.132634'
+long_version = '1.20240826.132634+git.8b552f5'

{returnn-1.20240824.1611 → returnn-1.20240826.132634}/returnn/frontend/array_.py

@@ -528,14 +528,45 @@ def masked_select(
 
     :param tensor:
     :param mask:
-    :param dims: the order of the dims defines the format. those dims should be exactly the dims of the mask.
-    :param out_dim:
-    :return: tensor where all dims in mask/dims are removed and replaced by a new dim.
-        the new dim is also returned.
-        if mask==True for all elements, the returned tensor would be simply the flattened input tensor.
-    """
-    # noinspection PyProtectedMember
-    return tensor._raw_backend.masked_select(tensor, mask=mask, dims=dims, out_dim=out_dim)
+    :param dims: The same or a subset of the dims of ``mask``.
+        The order of the dims defines the packed format (in ``out_dim``).
+        If only a single dim is given, only this dim is packed, while all other dims are kept as is.
+    :param out_dim: packed ``dims`` into this dim. Automatically created if not given, and then returned.
+    :return: Tensor where all ``dims`` are removed and replaced by a new dim (``out_dim``).
+        The new dim (``out_dim``) is also returned.
+        If mask==True for all elements, the returned tensor would be simply the flattened input tensor.
+    """
+    mask_dims_set = set(mask.dims)
+    dims_set = set(dims)
+    if not dims_set.issubset(mask_dims_set):
+        raise ValueError(f"masked_select: dims {dims} not subset of mask dims {mask.dims}")
+    if not dims_set:
+        raise ValueError(f"masked_select: dims {dims} empty")
+    if dims_set == mask_dims_set:
+        # noinspection PyProtectedMember
+        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.
+    if len(dims) != 1:
+        # Please check https://github.com/rwth-i6/returnn/issues/1605 when you need this.
+        raise NotImplementedError(f"masked_select: dims {dims} len {len(dims)} > 1 not supported")
+    (in_dim,) = dims
+    in_dim: Dim
+    mask = mask.copy_masked(mask_value=False, dims=dims)
+    idxs = rf.cumsum(rf.cast(mask, "int32"), spatial_dim=in_dim)  # [T,B] -> idx in T' + 1
+    new_size = rf.gather(idxs, indices=in_dim.get_size_tensor() - 1, axis=in_dim)  # [B]
+    if out_dim is None:
+        out_dim = Dim(new_size, name="masked_select")
+    elif out_dim.dyn_size_ext is None:
+        out_dim.dyn_size_ext = new_size
+    elif out_dim.dyn_size_ext is not None and out_dim.dyn_size_ext.raw_tensor is None:
+        out_dim.dyn_size_ext.raw_tensor = new_size.raw_tensor
+    new_time = rf.reduce_max(new_size, axis=new_size.dims)  # T'
+    idxs = rf.where(mask, idxs - 1, new_time)  # new_time is the padding idx
+    ext_out_dim = out_dim + 1  # one more for the padded data
+    ext_res = rf.scatter(tensor, indices=idxs, indices_dim=in_dim, out_dim=ext_out_dim)
+    res, _ = rf.slice(ext_res, axis=ext_out_dim, size=out_dim)
+    return res, out_dim
 
 
 def masked_scatter(source: Tensor, *, mask: Tensor, dims: Sequence[Dim], in_dim: Dim) -> Tensor:

{returnn-1.20240824.1611 → returnn-1.20240826.132634}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240824.1611 → returnn-1.20240826.132634}/tests/test_rf_array.py

@@ -33,6 +33,76 @@ def test_pack_padded():
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
 
 
+def test_masked_select():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim, in_dim], dtype="float32"),
+            "mask": Tensor("mask", [batch_dim, time_dim], dtype="bool"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        x = extern_data["data"]
+        mask = extern_data["mask"]
+        out, pack_dim = rf.masked_select(x, mask=mask, dims=[batch_dim, time_dim])
+        out.mark_as_default_output(shape=(pack_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
+def test_masked_select_single_dim():
+    # https://github.com/rwth-i6/returnn/issues/1605
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim, in_dim], dtype="float32"),
+            "mask": Tensor("mask", [batch_dim, time_dim], dtype="bool"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        x = extern_data["data"]
+        mask = extern_data["mask"]
+        out, pack_dim = rf.masked_select(x, mask=mask, dims=[time_dim])
+        out.mark_as_default_output(shape=(batch_dim, pack_dim, in_dim))
+        pack_dim.dyn_size_ext.mark_as_output("out_size", shape=[batch_dim])
+        x.mark_as_output("input", shape=[batch_dim, time_dim, in_dim])
+        mask.mark_as_output("mask", shape=[batch_dim, time_dim])
+        time_dim.dyn_size_ext.mark_as_output("in_size", shape=[batch_dim])
+
+    res = run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step, test_tensorflow=False)
+    for k, v in res.data.items():
+        print(f"{k} {v} = {v.raw_tensor}")
+    num_batches = res["input"].raw_tensor.shape[0]
+    have_non_zero = False
+    have_less = False
+    for b in range(num_batches):
+        in_size = res["in_size"].raw_tensor[b]
+        out_size = res["out_size"].raw_tensor[b]
+        print(f"batch {b}, in_size {in_size}, out_size {out_size}")
+        assert 0 <= out_size <= in_size
+        if out_size > 0:
+            have_non_zero = True
+        if out_size < in_size:
+            have_less = True
+        c = 0
+        for t in range(in_size):
+            if res["mask"].raw_tensor[b, t]:
+                c += 1
+        assert c == out_size
+        t_ = 0
+        for t in range(in_size):
+            if res["mask"].raw_tensor[b, t]:
+                assert (res["input"].raw_tensor[b, t] == res["output"].raw_tensor[b, t_]).all()
+                t_ += 1
+    assert have_non_zero and have_less  # just that the test case covered all cases
+
+
 def test_reshape():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

returnn-1.20240824.1611/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240824.001611'
-long_version = '1.20240824.001611+git.a15f05e'