returnn 1.20230418.124036__tar.gz → 1.20230418.134656__tar.gz

{returnn-1.20230418.124036/returnn.egg-info → returnn-1.20230418.134656}/PKG-INFO

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

returnn-1.20230418.134656/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230418.134656'
+long_version = '1.20230418.134656+git.79116cf'

{returnn-1.20230418.124036 → returnn-1.20230418.134656}/returnn/tf/frontend_layers/_backend.py

@@ -827,6 +827,12 @@ class ReturnnLayersBackend(Backend[Layer]):
                 dilation_rate=dilation_rate,
                 padding=padding,
             )
+        other_dims = [d for d in source.dims if d not in in_spatial_dims and not d.is_batch_dim()]
+        assert other_dims  # currently not implemented otherwise, need in_dim...
+        if source.feature_dim and source.feature_dim in other_dims:
+            in_dim = source.feature_dim
+        else:
+            in_dim = other_dims[-1]
         args = {
             "mode": mode,
             "pool_size": pool_size,
@@ -835,8 +841,15 @@ class ReturnnLayersBackend(Backend[Layer]):
             "strides": strides,
             "in_spatial_dims": in_spatial_dims,
             "out_spatial_dims": out_spatial_dims,
+            "in_dim": in_dim,  # it does not really matter, but we need sth currently
         }
         layer = rfl.make_layer({"class": "pool", "from": source, **args}, name="pool")
+        if source.feature_dim != in_dim:
+            # We want that the feature-dim stays consistent. PoolLayer currently just sets it to the in_dim.
+            layer = rfl.make_layer(
+                {"class": "reinterpret_data", "from": layer, "set_axes": {"F": source.feature_dim}},
+                name="pool_reset_feature",
+            )
         return layer, out_spatial_dims
 
 

{returnn-1.20230418.124036 → returnn-1.20230418.134656}/returnn/tf/layers/basic.py

@@ -5591,7 +5591,7 @@ class ReinterpretDataLayer(_ConcatInputLayer):
         :param str|list[str] switch_axes: e.g. "bt" to switch batch and time axes
         :param LayerBase|None size_base: copy the size_placeholder from the given layer
         :param LayerBase|None batch_dim_base: copy the batch dim from this layer
-        :param dict[str,Dim|str] set_axes:
+        :param dict[str,Dim|str|None] set_axes:
           This can be used to overwrite the special axes like time_dim_axis or feature_dim_axis.
           For that, use keys "B","T" or "F", and a value via :func:`Data.get_axis_from_description`.
         :param dict[str|Dim,Dim]|None set_dim_tags: axis -> new dim tag. assigns new dim tags.
@@ -5698,7 +5698,7 @@ class ReinterpretDataLayer(_ConcatInputLayer):
         :param str|list[str] switch_axes: e.g. "bt" to switch batch and time axes
         :param LayerBase|None size_base: similar as size_target
         :param LayerBase|None batch_dim_base:
-        :param dict[str,Dim|str] set_axes:
+        :param dict[str,Dim|str|None] set_axes:
         :param dict[str|Dim,Dim]|None set_dim_tags:
         :param bool enforce_batch_major:
         :param bool enforce_time_major:
@@ -5738,10 +5738,9 @@ class ReinterpretDataLayer(_ConcatInputLayer):
                 s = map_axis_name(s)
                 if isinstance(i, int):
                     assert enforce_batch_major or enforce_time_major, "%r: explicit set_axes %r" % (name, set_axes)
-                i = out.get_axis_from_description(i)
+                if i is not None:
+                    i = out.get_axis_from_description(i)
                 setattr(out, s, i)
-                if s == "feature_dim_axis":
-                    out.dim = out.batch_shape[out.feature_dim_axis]
         if out.size_placeholder and size_base:  # size_placeholder might be None, e.g. via DataNotAvailableLayer
             assert size_base.output.size_placeholder
             assert len(out.size_placeholder) == len(size_base.output.size_placeholder)
@@ -6517,7 +6516,7 @@ class PoolLayer(_ConcatInputLayer):
         elif isinstance(strides, int):
             strides = [strides] * len(pool_size)
         assert len(strides) == len(pool_size)
-        super(PoolLayer, self).__init__(**kwargs)
+        super(PoolLayer, self).__init__(in_dim=in_dim, out_dim=out_dim, **kwargs)
         assert not self.input_data.sparse
         assert self.input_data.have_batch_axis()
         assert (

{returnn-1.20230418.124036 → returnn-1.20230418.134656}/returnn/torch/frontend/_backend.py

@@ -515,6 +515,8 @@ class TorchBackend(Backend[torch.Tensor]):
         :return: a `kind` b
         """
         assert a.dim() == b.dim()
+        if kind == "squared_difference":
+            return (a - b) ** 2
         kind = {
             "truediv": "true_divide",
             "floordiv": "floor_divide",
@@ -1091,7 +1093,6 @@ class TorchBackend(Backend[torch.Tensor]):
                 dilation_rate=dilation_rate,
                 padding=padding,
             )
-        assert padding == "valid"  # not implemented otherwise
         batch_dims = [d for d in source.dims if d not in tuple(in_spatial_dims)]
         # Torch conv expects (N,C,<spatial dims>) as shape.
         # batch_dims would actually cover the channel-dim (C) as well,
@@ -1101,8 +1102,9 @@ class TorchBackend(Backend[torch.Tensor]):
             source.raw_tensor,
             # Potentially merge batch dims all together.
             # Keep the last as the channel-dim, but not sure if this is really relevant.
-            [-1, batch_dims[-1] if batch_dims else 1] + [d.get_dim_value() for d in in_spatial_dims],
+            [-1, batch_dims[-1].get_dim_value() if batch_dims else 1] + [d.get_dim_value() for d in in_spatial_dims],
         )
+        ceil_mode = padding.lower() == "same"
         if len(in_spatial_dims) == 1:
             # There is also conv_tbc, but it's a bit limited (no dilation)
             # and also unclear when exactly it is faster.
@@ -1111,6 +1113,7 @@ class TorchBackend(Backend[torch.Tensor]):
                 kernel_size=pool_size,
                 stride=strides,
                 dilation=dilation_rate or 1,
+                ceil_mode=ceil_mode,
             )
         elif len(in_spatial_dims) == 2:
             out_raw = torch.nn.functional.max_pool2d(
@@ -1118,6 +1121,7 @@ class TorchBackend(Backend[torch.Tensor]):
                 kernel_size=pool_size,
                 stride=strides,
                 dilation=dilation_rate or 1,
+                ceil_mode=ceil_mode,
             )
         elif len(in_spatial_dims) == 3:
             out_raw = torch.nn.functional.max_pool3d(
@@ -1125,6 +1129,7 @@ class TorchBackend(Backend[torch.Tensor]):
                 kernel_size=pool_size,
                 stride=strides,
                 dilation=dilation_rate or 1,
+                ceil_mode=ceil_mode,
             )
         else:
             raise ValueError(f"invalid number of filter dims {in_spatial_dims}, expected 1, 2, or 3")

{returnn-1.20230418.124036 → returnn-1.20230418.134656/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230418.124036 → returnn-1.20230418.134656}/tests/test_rf_conv.py

@@ -69,3 +69,45 @@ def test_conv1d_depthwise():
         out.mark_as_default_output(shape=(batch_dim, spatial_dim, out_dim))
 
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
+
+
+def test_maxpool1d_padding_valid():
+    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"),
+        }
+    )
+
+    class _Net(rf.Module):
+        def __call__(self, x: rf.Tensor, *, in_spatial_dim: Dim) -> Tuple[Tensor, Dim]:
+            return rf.max_pool1d(x, pool_size=3, padding="valid", in_spatial_dim=in_spatial_dim)
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, model: _Net, extern_data: TensorDict):
+        out, out_spatial_dim = model(extern_data["data"], in_spatial_dim=time_dim)
+        out.mark_as_default_output(shape=(batch_dim, out_spatial_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
+
+
+def test_maxpool1d_padding_same():
+    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"),
+        }
+    )
+
+    class _Net(rf.Module):
+        def __call__(self, x: rf.Tensor, *, in_spatial_dim: Dim) -> Tuple[Tensor, Dim]:
+            return rf.max_pool1d(x, pool_size=3, padding="same", in_spatial_dim=in_spatial_dim)
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, model: _Net, extern_data: TensorDict):
+        out, out_spatial_dim = model(extern_data["data"], in_spatial_dim=time_dim)
+        out.mark_as_default_output(shape=(batch_dim, out_spatial_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)

{returnn-1.20230418.124036 → returnn-1.20230418.134656}/tests/test_rf_math.py

@@ -28,3 +28,25 @@ def test_neg():
         out.mark_as_default_output(shape=(batch_dim, time_dim, in_dim))
 
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
+
+
+def test_squared_difference():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "a": Tensor("a", [batch_dim, time_dim, in_dim], dtype="float32"),
+            "b": Tensor("b", [batch_dim, time_dim, in_dim], dtype="float32"),
+        }
+    )
+
+    class _Net(rf.Module):
+        def __call__(self, a: Tensor, b: Tensor) -> Tensor:
+            return rf.squared_difference(a, b)
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, model: _Net, extern_data: TensorDict):
+        out = model(extern_data["a"], extern_data["b"])
+        out.mark_as_default_output(shape=(batch_dim, time_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)

returnn-1.20230418.124036/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230418.124036'
-long_version = '1.20230418.124036+git.2ec964a'