returnn 1.20250119.201443__tar.gz → 1.20250120.153919__tar.gz

{returnn-1.20250119.201443/returnn.egg-info → returnn-1.20250120.153919}/PKG-INFO

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

returnn-1.20250120.153919/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250120.153919'
+long_version = '1.20250120.153919+git.d86ab9f'

{returnn-1.20250119.201443 → returnn-1.20250120.153919}/returnn/torch/frontend/_backend.py

@@ -674,6 +674,12 @@ class TorchBackend(Backend[torch.Tensor]):
         assert targets.sparse_dim and targets.sparse_dim.dimension <= logits.feature_dim.dimension
         # PyTorch expects the logits to be of shape (T, B, C) where T is the input spatial dim.
         batch_dims = logits.remaining_dims((input_spatial_dim, logits.feature_dim))
+        batch_dims_targets = targets.remaining_dims(targets_spatial_dim)
+        if set(batch_dims) != set(batch_dims_targets):
+            # Need to broadcast.
+            logits = rf.expand_dims(logits, [d for d in batch_dims_targets if d not in batch_dims])
+            targets = rf.expand_dims(targets, [d for d in batch_dims if d not in batch_dims_targets])
+            batch_dims = logits.remaining_dims((input_spatial_dim, logits.feature_dim))
         batch_shape = [d.get_dim_value() for d in batch_dims]
         batch_n_elems = prod(batch_shape)
         logits = logits.copy_transpose([input_spatial_dim] + batch_dims + [logits.feature_dim])
@@ -1165,6 +1171,9 @@ class TorchBackend(Backend[torch.Tensor]):
             start = 0
         if isinstance(size, Dim):
             size = size.get_dim_value()
+        elif isinstance(size, Tensor):
+            assert size.dims == ()  # scalar
+            size = size.raw_tensor
         if size is not None:
             assert end is None
             out.raw_tensor = torch.narrow(source.raw_tensor, dim=axis_int, start=start, length=size)

{returnn-1.20250119.201443 → returnn-1.20250120.153919/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250119.201443 → returnn-1.20250120.153919}/tests/test_rf_array.py

@@ -542,6 +542,26 @@ def test_slice():
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
 
 
+def test_slice_dyn_size():
+    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"]
+        size = time_dim.get_size_tensor()
+        size = rf.reduce_min(size, axis=size.dims)  # do whatever, but get some scalar
+        pack, new_time = rf.slice(x, axis=time_dim, size=size)
+        pack.mark_as_default_output(shape=(batch_dim, new_time, in_dim))
+
+    run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step)
+
+
 def test_shift_right():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

returnn-1.20250119.201443/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250119.201443'
-long_version = '1.20250119.201443+git.d4343d7'