returnn 1.20240926.134856__tar.gz → 1.20241001.90825__tar.gz

{returnn-1.20240926.134856 → returnn-1.20241001.90825}/PKG-INFO

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

returnn-1.20241001.90825/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241001.090825'
+long_version = '1.20241001.090825+git.8309a61'

{returnn-1.20240926.134856 → returnn-1.20241001.90825}/returnn/frontend/_backend.py

@@ -992,6 +992,16 @@ class Backend(Generic[T]):
         """is finite"""
         raise NotImplementedError
 
+    @staticmethod
+    def is_infinite(x: Tensor) -> Tensor:
+        """is positive or negative infinite"""
+        raise NotImplementedError
+
+    @staticmethod
+    def is_neg_infinite(x: Tensor) -> Tensor:
+        """is negative infinite"""
+        raise NotImplementedError
+
     @staticmethod
     def clip_by_value(
         x: Tensor,

{returnn-1.20240926.134856 → returnn-1.20241001.90825}/returnn/frontend/array_.py

@@ -41,6 +41,7 @@ __all__ = [
     "scatter_argmax",
     "slice",
     "shift_right",
+    "shift_left",
     "reverse_sequence",
     "where",
     "search_sorted",
@@ -893,6 +894,15 @@ def shift_right(source: Tensor, *, axis: Dim, pad_value: Union[rf.RawTensorTypes
     return padded_slice
 
 
+def shift_left(source: Tensor, *, axis: Dim, pad_value: Union[rf.RawTensorTypes, Tensor], amount: int = 1) -> Tensor:
+    """shift left by amount, pad right with right_pad"""
+    padded, (padded_dim,) = rf.pad(
+        source, axes=[axis], padding=[(0, amount)], mode="constant", value=pad_value, handle_dynamic_dims=True
+    )
+    padded_slice, _ = rf.slice(padded, axis=padded_dim, start=amount, size=axis)
+    return padded_slice
+
+
 def reverse_sequence(tensor: Tensor, *, axis: Dim, handle_dynamic_dims: bool = True) -> Tensor:
     """
     Similar as tf.reverse_sequence, or Torch flip (but taking seq lengths into account).

{returnn-1.20240926.134856 → returnn-1.20241001.90825}/returnn/frontend/math_.py

@@ -38,6 +38,8 @@ __all__ = [
     "opt_logical_or",
     "opt_logical_and",
     "is_finite",
+    "is_infinite",
+    "is_neg_infinite",
     "maximum",
     "minimum",
     "clip_by_value",
@@ -368,6 +370,18 @@ def is_finite(a: Tensor) -> Tensor:
     return a._raw_backend.is_finite(a)
 
 
+def is_infinite(a: Tensor) -> Tensor:
+    """is positive or negative infinite"""
+    # noinspection PyProtectedMember
+    return a._raw_backend.is_infinite(a)
+
+
+def is_neg_infinite(a: Tensor) -> Tensor:
+    """is negative infinite"""
+    # noinspection PyProtectedMember
+    return a._raw_backend.is_neg_infinite(a)
+
+
 def maximum(a: Tensor, b: Union[Tensor, _RawTensorTypes], *other_tensors) -> Tensor:
     """maximum"""
     if not other_tensors:

{returnn-1.20240926.134856 → returnn-1.20241001.90825}/returnn/tensor/_dim_extra.py

@@ -1643,25 +1643,14 @@ class _DimMixin:
         if self._extra:
             self._extra.derived_from_op = None
             self._extra.derived_from_tag = None
-        if self_same_as is not self:
-            assert not self_same_as.same_as
-            if self_same_as is other_same_base:
-                return
-            with util.guard_infinite_recursion(_d.Dim.declare_same_as, self_same_as, other_same_base):
-                self_same_as.declare_same_as(other_same_base)
-            if (self.dyn_size_ext is None or not self._validate_in_current_graph()) and self_same_as.dyn_size_ext:
-                self.dyn_size_ext = self_same_as.get_dyn_size_ext_for_batch_ctx(
-                    self.batch, self.control_flow_ctx, template_only=True
-                )
-        else:
-            if self._extra:
-                for dim_ in self._extra.same_for_batch_ctx.values():
+        if self._extra:
+            for dim_ in self._extra.same_for_batch_ctx.values():
+                # noinspection PyProtectedMember
+                if dim_._extra:
                     # noinspection PyProtectedMember
-                    if dim_._extra:
-                        # noinspection PyProtectedMember
-                        dim_._extra.derived_from_op = None
-                        # noinspection PyProtectedMember
-                        dim_._extra.derived_from_tag = None
+                    dim_._extra.derived_from_op = None
+                    # noinspection PyProtectedMember
+                    dim_._extra.derived_from_tag = None
         # Now merge existing variants. But only if not derived via op, because in that case, we can (and should!)
         # automatically infer it. Note that we only got here when the other is not the same dim, so it means that
         # the other is really different, the sizes are potentially different, but we want to overtake the other.
@@ -1723,6 +1712,13 @@ class _DimMixin:
                 other_same_base.batch, other_same_base.control_flow_ctx, template_only=True
             )
             other_same_base._maybe_update()
+        if (
+            self.dyn_size_ext
+            and self.dyn_size_ext.raw_tensor is None
+            and other_same_base.dyn_size_ext.raw_tensor is not None
+        ):
+            self.dyn_size_ext = other_same_base.dyn_size_ext.copy()
+            self._maybe_update()
         if self.is_dim_known() and other.is_dim_known():
             assert self.dimension == other.dimension
         elif self.is_dim_known() and not other.is_dim_known():
@@ -1762,6 +1758,13 @@ class _DimMixin:
             if self_ is not self:
                 self.control_flow_ctx = self_.control_flow_ctx  # might be different
                 self.dyn_size_ext = self_.dyn_size_ext  # might be unset
+        if self_same_as is not self:
+            assert not self_same_as.same_as
+            assert self_same_as is not other_same_base
+            with util.guard_infinite_recursion(_d.Dim.declare_same_as, self_same_as, other_same_base):
+                # Do this at the end, when this is not our own base anymore,
+                # because otherwise any reset() might incorrectly reset the new other base.
+                self_same_as.declare_same_as(other_same_base)
 
     def _merge_same_for_batch_ctx_dict(self: _d.Dim, other: _d.Dim):
         """

{returnn-1.20240926.134856 → returnn-1.20241001.90825}/returnn/torch/frontend/_backend.py

@@ -954,7 +954,9 @@ class TorchBackend(Backend[torch.Tensor]):
             raise TypeError(f"Unsupported type for indices: {type(indices)}")
         if clip_to_valid:
             if axis.dyn_size_ext:
-                indices = rf.clip_by_value(indices, 0, axis.get_dyn_size_ext_for_device(indices.device) - 1)
+                indices = rf.clip_by_value(
+                    indices, 0, axis.get_dyn_size_ext_for_device(indices.device) - 1, allow_broadcast_all_sources=True
+                )
             else:
                 indices = indices.copy()
                 indices.raw_tensor = torch.clamp(indices.raw_tensor, 0, source.raw_tensor.shape[axis_int] - 1)
@@ -1212,6 +1214,20 @@ class TorchBackend(Backend[torch.Tensor]):
         out.raw_tensor = torch.isfinite(x.raw_tensor)
         return out
 
+    @staticmethod
+    def is_infinite(x: Tensor) -> Tensor:
+        """is positive or negative infinite"""
+        out = x.copy_template("is_infinite", dtype="bool")
+        out.raw_tensor = torch.isinf(x.raw_tensor)
+        return out
+
+    @staticmethod
+    def is_neg_infinite(x: Tensor) -> Tensor:
+        """is negative infinite"""
+        out = x.copy_template("is_neg_infinite", dtype="bool")
+        out.raw_tensor = torch.isneginf(x.raw_tensor)
+        return out
+
     @staticmethod
     def clip_by_value(
         x: Tensor,
@@ -1404,7 +1420,7 @@ class TorchBackend(Backend[torch.Tensor]):
         if use_mask and any(dim.need_masking() for dim in axis):
             source = source.copy()
             dtype = source.raw_tensor.dtype
-            if mode == "max":
+            if mode in ("max", "logsumexp"):
                 mask_value = torch.finfo(dtype).min if dtype.is_floating_point else torch.iinfo(dtype).min
             elif mode == "min":
                 mask_value = torch.finfo(dtype).max if dtype.is_floating_point else torch.iinfo(dtype).max

{returnn-1.20240926.134856 → returnn-1.20241001.90825}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240926.134856 → returnn-1.20241001.90825}/tests/test_rf_base.py

@@ -263,6 +263,30 @@ def test_dim_value():
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
 
 
+def test_dim_size_after_declare_same_as():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim], dtype="float32"),
+        }
+    )
+    out_spatial_dim = Dim(None, name="out_spatial")
+
+    # noinspection PyShadowingNames
+    def _forward_step(**_kwargs):
+        assert time_dim.dyn_size_ext.raw_tensor is not None
+        enc_spatial_dim = time_dim * 2
+        assert enc_spatial_dim.dyn_size_ext.raw_tensor is not None
+        out_spatial_dim.declare_same_as(enc_spatial_dim)
+        assert enc_spatial_dim.dyn_size_ext.raw_tensor is not None
+        assert out_spatial_dim.dyn_size_ext.raw_tensor is not None
+        out_spatial_dim.dyn_size_ext.mark_as_default_output(shape=[batch_dim])
+
+    # Running it twice can trigger some weird behavior in Dim.declare_same_as.
+    run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+    run_model(extern_data, lambda **_kwargs: rf.Module(), _forward_step, test_tensorflow=False)
+
+
 def test_dim_mask():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

returnn-1.20240926.134856/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240926.134856'
-long_version = '1.20240926.134856+git.1fcaadd'