returnn 1.20250117.110003__tar.gz → 1.20250117.155146__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250117.110003
+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.20250117.110003 → 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.20250117.110003 → 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.20250117.110003 → returnn-1.20250117.155146/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250117.110003
+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.110003 → returnn-1.20250117.155146}/tests/test_rf_array.py

@@ -48,11 +48,31 @@ def test_pack_padded_md():
         labels = extern_data["labels"]
         assert hyps_spatial_dim in labels.dims
         assert hyps_spatial_dim.dyn_size_ext.dims == (batch_dim, hyps_dim)
+        labels.mark_as_output("labels", shape=(batch_dim, hyps_dim, hyps_spatial_dim))
+        hyps_spatial_dim.dyn_size_ext.mark_as_output("hyps_spatial_size", shape=[batch_dim, hyps_dim])
         labels_, _ = rf.pack_padded(labels, dims=[hyps_dim, hyps_spatial_dim], out_dim=hyps_packed_spatial_dim)
         assert hyps_packed_spatial_dim.dyn_size_ext.dims == (batch_dim,)
+        hyps_packed_spatial_dim.dyn_size_ext.mark_as_output("packed_size", shape=[batch_dim])
         labels_.mark_as_default_output(shape=[batch_dim, hyps_packed_spatial_dim])
 
-    run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+    out = run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+    print("out:")
+    print(out["output"].raw_tensor.shape)
+    print(out["output"].raw_tensor)
+    batch_size = out["labels"].raw_tensor.shape[0]
+    for b in range(batch_size):
+        num_hyps = hyps_dim.dimension
+        hyps_spatial_sizes = out["hyps_spatial_size"].raw_tensor[b]  # [hyps]
+        hyps_packed_spatial_size = out["packed_size"].raw_tensor[b]  # []
+        print(f"batch {b}, hyps_spatial_size {hyps_spatial_sizes}, hyps_packed_spatial_size {hyps_packed_spatial_size}")
+        assert hyps_packed_spatial_size == sum(hyps_spatial_sizes)
+        labels = [out["labels"].raw_tensor[b, h, : hyps_spatial_sizes[h]] for h in range(num_hyps)]
+        labels_ = [seq.tolist() for seq in labels]
+        print("labels:", labels_)
+        packed = out["output"].raw_tensor[b, :hyps_packed_spatial_size]  # [hyps_packed_spatial]
+        packed_ = packed.tolist()
+        print("packed:", packed_)
+        assert sum(labels_, []) == packed_
 
 
 def test_masked_select():
@@ -125,6 +145,61 @@ def test_masked_select_single_dim():
     assert have_non_zero and have_less  # just that the test case covered all cases
 
 
+def test_pad_packed():
+    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"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        x = extern_data["data"]
+        x.mark_as_output("in", shape=(batch_dim, time_dim, in_dim))
+        x_, batch_time_flat_dim = rf.pack_padded(x, dims=[batch_dim, time_dim])
+        assert x_.dims_set == {batch_time_flat_dim, in_dim}
+        # test [batch_fime_flat,feat], batch_fime_flat -> [batch, time]
+        x__ = rf.pad_packed(x_, in_dim=batch_time_flat_dim, dims=[batch_dim, time_dim])
+        assert x__.dims_set == {batch_dim, time_dim, in_dim}
+        x__.mark_as_output("out", shape=(batch_dim, time_dim, in_dim))
+
+    out_dict = run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+    in_ = out_dict["in"]
+    out = out_dict["out"]
+    np.testing.assert_array_equal(in_.raw_tensor, out.raw_tensor)
+
+
+def test_pad_packed_batched():
+    batch_dim_ = Dim(3, name="batch")
+    in_dim = Dim(2, name="in")
+
+    # noinspection PyShadowingNames
+    def _forward_step(**_kwargs):
+        time1_dim = Dim(rf.convert_to_tensor(np.array([1, 3, 4]), name="time1", dims=[batch_dim_], dtype="int32"))
+        time2_dim = Dim(rf.convert_to_tensor(np.array([3, 2, 2]), name="time2", dims=[batch_dim_], dtype="int32"))
+        x = rf.random_uniform((batch_dim_, time1_dim, time2_dim, in_dim))
+        x.mark_as_output("in", shape=(batch_dim_, time1_dim, time2_dim, in_dim))
+        x_, time_flat_dim = rf.pack_padded(x, dims=[time1_dim, time2_dim])
+        assert x_.dims_set == {batch_dim_, time_flat_dim, in_dim}
+        x_.mark_as_output("flat", shape=(batch_dim_, time_flat_dim, in_dim))
+        # test [batch,mult_fime_flat,feat], mult_fime_flat -> [time1, time2]
+        x__ = rf.pad_packed(x_, in_dim=time_flat_dim, dims=[time1_dim, time2_dim])
+        assert x__.dims_set == {batch_dim_, time1_dim, time2_dim, in_dim}
+        x__.mark_as_output("out", shape=(batch_dim_, time1_dim, time2_dim, in_dim))
+
+    out_dict = run_model(TensorDict(), lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+    in_ = out_dict["in"]
+    flat = out_dict["flat"]
+    out = out_dict["out"]
+    print(in_.raw_tensor.shape, flat.raw_tensor.shape, out.raw_tensor.shape)
+    print(in_.raw_tensor)
+    print(flat.raw_tensor)
+    print(out.raw_tensor)
+    np.testing.assert_array_equal(in_.raw_tensor, out.raw_tensor)
+
+
 def test_reshape():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

returnn-1.20250117.110003/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250117.110003'
-long_version = '1.20250117.110003+git.44fde7a'