returnn 1.20240917.154138__tar.gz → 1.20240918.101247__tar.gz

{returnn-1.20240917.154138 → returnn-1.20240918.101247}/PKG-INFO

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

returnn-1.20240918.101247/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240918.101247'
+long_version = '1.20240918.101247+git.2a3a1f8'

{returnn-1.20240917.154138 → returnn-1.20240918.101247}/returnn/frontend/array_.py

@@ -549,14 +549,16 @@ 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.
-    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)
+    if len(dims) > 1:
+        # Flatten it, in the specified order.
+        tensor, in_dim = rf.merge_dims(tensor, dims=dims)
+        mask, _ = rf.merge_dims(mask, dims=dims, out_dim=in_dim)
+    else:
+        (in_dim,) = dims
+    in_dim: Dim
     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]
+    new_size = rf.gather(idxs, indices=in_dim.get_dim_value_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:
@@ -620,7 +622,6 @@ def pack_padded(
     """
     assert not enforce_sorted  # not implemented yet...
     mask = rf.sequence_mask(dims, device=source.device)
-    assert mask.dims_set == set(dims)
     # Note: We could already calculate out_dim here, as follows:
     #   out_dim = Dim(rf.num_elements_of_shape(dims), name="packed")
     # This could trigger a more efficient calculation path in masked_select,

{returnn-1.20240917.154138 → returnn-1.20240918.101247}/returnn/tf/layers/basic.py

@@ -4589,7 +4589,6 @@ class MergeDimsLayer(_ConcatInputLayer):
             or input_data.feature_dim_axis_or_unspecified is not NotSpecified
         ):
             data.feature_dim_axis = new_feature_dim_axis  # explicitly set
-            data.dim = data.batch_shape[data.feature_dim_axis] if data.feature_dim_axis is not None else None
 
         if input_data.batch_dim_axis in axes and data.batch:
             for axis in axes:

{returnn-1.20240917.154138 → returnn-1.20240918.101247}/returnn/torch/engine.py

@@ -1086,6 +1086,9 @@ class Engine(EngineBase):
             step_idx = 0
             for extern_data_raw in data_loader:
                 step_begin_time = time.time()
+                if self._forward_step_expected_outputs:
+                    # Also resets any dyn dims, which might have been set in the prev step.
+                    self._forward_step_expected_outputs.reset_content()
                 extern_data = extern_data_util.raw_dict_to_extern_data(
                     extern_data_raw, extern_data_template=self.extern_data, device=self._device
                 )

{returnn-1.20240917.154138 → returnn-1.20240918.101247}/returnn/torch/frontend/_backend.py

@@ -980,7 +980,12 @@ class TorchBackend(Backend[torch.Tensor]):
             index_ext_dims = list(source.dims)
             index_ext_dims[axis_int] = index_own_dims_flat
             assert indices.dims == tuple(index_ext_dims)
-            out_raw = torch.gather(source.raw_tensor, dim=axis_int, index=indices.raw_tensor.type(torch.int64))
+            out_raw = torch.gather(
+                source.raw_tensor,
+                dim=axis_int,
+                # torch.gather wants indices as int64, and the indices should be the same device as the source.
+                index=indices.raw_tensor.type(torch.int64).to(source.raw_tensor.device),
+            )
             if len(index_own_dims) == 1:
                 pass  # nothing to do
             elif len(index_own_dims) == 0:

{returnn-1.20240917.154138 → returnn-1.20240918.101247}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240917.154138 → returnn-1.20240918.101247}/tests/test_rf_array.py

@@ -33,6 +33,28 @@ def test_pack_padded():
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
 
 
+def test_pack_padded_md():
+    batch_dim = Dim(3, name="batch")
+    hyps_dim = Dim(5, name="hyps")
+    hyps_spatial_dim = Dim(Tensor("hyps_spatial", [batch_dim, hyps_dim], dtype="int32"))
+    vocab_dim = Dim(1000, name="vocab")
+    extern_data = TensorDict(
+        {"labels": Tensor("labels", [batch_dim, hyps_dim, hyps_spatial_dim], dtype="int32", sparse_dim=vocab_dim)}
+    )
+    hyps_packed_spatial_dim = Dim(None, name="hyps_packed_spatial")
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        labels = extern_data["labels"]
+        assert hyps_spatial_dim in labels.dims
+        assert hyps_spatial_dim.dyn_size_ext.dims == (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,)
+        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)
+
+
 def test_masked_select():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

returnn-1.20240917.154138/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240917.154138'
-long_version = '1.20240917.154138+git.627ab49'