returnn 1.20230519.174327__tar.gz → 1.20230521.10102__tar.gz

{returnn-1.20230519.174327/returnn.egg-info → returnn-1.20230521.10102}/PKG-INFO

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

returnn-1.20230521.10102/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230521.010102'
+long_version = '1.20230521.010102+git.bb22c48'

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/demos/demo-rf.config

@@ -26,7 +26,7 @@ extern_data = {
     "classes": {"dims": (batch_dim, time_dim), "sparse_dim": out_dim, "dtype": "int32"},
 }
 model_outputs = {
-    "output": {"dim": out_dim.size},
+    "output": {"dims": (batch_dim, time_dim, out_dim), "dtype": "float32"},
 }
 
 batching = "random"

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/demos/demo-torch.config

@@ -38,17 +38,13 @@ max_seqs = 10
 class Model(torch.nn.Module):
     def __init__(self):
         super().__init__()
-        kernel_size = 5
-        # We manually calculate the padding because padding="same"
-        # is not exportable to ONNX as of 2023/05/19.
-        # When stride = 1 and kernel_size % 2 = 1, the following formula applies.
-        padding = (kernel_size - 1) // 2
+        # Note that padding="same" is not exportable to ONNX as of 2023/05/19.
         self.layers = nn.Sequential(
-            nn.Conv1d(num_inputs, 50, kernel_size, padding=padding),
+            nn.Conv1d(num_inputs, 50, 5, padding=2),
             nn.ReLU(),
-            nn.Conv1d(50, 100, kernel_size, padding=padding),
+            nn.Conv1d(50, 100, 5, padding=2),
             nn.ReLU(),
-            nn.Conv1d(100, num_outputs, kernel_size, padding=padding),
+            nn.Conv1d(100, num_outputs, 5, padding=2),
         )
 
     def forward(self, x):  # (B, T, F)

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/returnn/frontend/_backend.py

@@ -773,6 +773,17 @@ class Backend(Generic[T]):
         """slice"""
         raise NotImplementedError
 
+    @staticmethod
+    def where(
+        cond: Tensor,
+        true_: Union[Tensor, rf.RawTensorTypes],
+        false_: Union[Tensor, rf.RawTensorTypes],
+        *,
+        allow_broadcast_all_sources: bool = False,
+    ) -> Tensor:
+        """where"""
+        raise NotImplementedError
+
     @staticmethod
     def matmul(a: Tensor[T], b: Tensor[T], *, reduce: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> Tensor[T]:
         """

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/returnn/frontend/array_.py

@@ -31,6 +31,9 @@ __all__ = [
     "gather",
     "slice",
     "shift_right",
+    "where",
+    "sparse_to_dense",
+    "one_hot",
 ]
 
 
@@ -490,3 +493,52 @@ def shift_right(source: Tensor, *, axis: Dim, pad_value: Union[rf.RawTensorTypes
     padded, (padded_dim,) = rf.pad(source, axes=[axis], padding=[(amount, 0)], mode="constant", value=pad_value)
     padded_slice, _ = rf.slice(padded, axis=padded_dim, size=axis)
     return padded_slice
+
+
+def where(
+    cond: Tensor,
+    true_: Union[Tensor, rf.RawTensorTypes],
+    false_: Union[Tensor, rf.RawTensorTypes],
+    *,
+    allow_broadcast_all_sources: bool = False,
+) -> Tensor:
+    """
+    Wraps tf.where, which is SwitchLayer in RETURNN.
+
+    :return: true_ if cond else false_, elemwise.
+    """
+    # noinspection PyProtectedMember
+    return cond._raw_backend.where(cond, true_, false_, allow_broadcast_all_sources=allow_broadcast_all_sources)
+
+
+def sparse_to_dense(
+    labels: Union[Tensor, rf.RawTensorTypes],
+    *,
+    label_value: Union[Tensor, rf.RawTensorTypes],
+    other_value: Union[Tensor, rf.RawTensorTypes],
+    axis: Optional[Dim] = None,
+) -> Tensor:
+    """
+    Converts a sparse tensor to a dense one.
+
+    This is a more generic variant of "one_hot".
+
+    Note that usually this is not needed as most other functions should handle sparse tensors just fine
+    and much more efficiently than they would be with dense tensors.
+    """
+    labels = rf.convert_to_tensor(labels)
+    if not axis:
+        assert labels.sparse_dim, "sparse_to_dense: either provide `axis` or `labels` with sparse_dim"
+        axis = labels.sparse_dim
+    indices = rf.range_over_dim(axis)
+    return where(rf.compare_bc(labels, "==", indices), label_value, other_value)
+
+
+def one_hot(source: Tensor) -> Tensor:
+    """
+    one_hot. special case of :func:`sparse_to_dense`.
+
+    Note that usually this is not needed as most other functions should handle sparse tensors just fine
+    and much more efficiently than they would be with dense tensors.
+    """
+    return sparse_to_dense(source, label_value=1.0, other_value=0.0)

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/returnn/frontend/math_.py

@@ -12,6 +12,7 @@ from .types import RawTensorTypes as _RawTensorTypes
 
 __all__ = [
     "compare",
+    "compare_bc",
     "combine",
     "combine_bc",
     "equal",
@@ -115,6 +116,17 @@ def compare(
     return backend.compare(a, kind, b, allow_broadcast_all_sources=allow_broadcast_all_sources, dim_order=dim_order)
 
 
+def compare_bc(
+    a: Tensor,
+    kind: str,
+    b: Tensor,
+    *,
+    dim_order: Optional[Sequence[Dim]] = None,
+) -> Tensor:
+    """:func:`compare` with allow_broadcast_all_sources=True"""
+    return compare(a, kind, b, allow_broadcast_all_sources=True, dim_order=dim_order)
+
+
 @typing.overload
 def combine(
     a: Tensor,

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/returnn/tensor/utils.py

@@ -101,6 +101,8 @@ def tensor_fill_random_numpy_(
             if x.sparse_dim and x.sparse_dim.dimension is not None:
                 max_val = x.sparse_dim.dimension
             x.raw_tensor = rnd.randint(min_val, max_val, size=shape, dtype=x.dtype)
+        elif x.dtype == "bool":
+            x.raw_tensor = rnd.randint(0, 2, size=shape, dtype=x.dtype)
         elif x.dtype.startswith("float"):
             x.raw_tensor = rnd.normal(0.0, 1.0, size=shape).astype(x.dtype)
         elif x.dtype.startswith("complex"):

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/returnn/tf/frontend_layers/_backend.py

@@ -558,6 +558,20 @@ class ReturnnLayersBackend(Backend[Layer]):
             {"class": "slice", "from": source, "axis": axis, "out_dim": out_dim, **args}, name="slice"
         )
 
+    @staticmethod
+    def where(
+        cond: Tensor,
+        true_: Union[Tensor, rf.RawTensorTypes],
+        false_: Union[Tensor, rf.RawTensorTypes],
+        *,
+        allow_broadcast_all_sources: bool = False,
+    ) -> Tensor:
+        """where"""
+        allow_broadcast_all_sources  # noqa # ignore allow_broadcast_all_sources for now..., not implemented
+        return rfl.make_layer(
+            {"class": "switch", "condition": cond, "true_from": true_, "false_from": false_}, name="where"
+        )
+
     @staticmethod
     def matmul(a: Tensor, b: Tensor, *, reduce: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> Tensor:
         """matmul"""

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/returnn/torch/frontend/_backend.py

@@ -626,6 +626,8 @@ class TorchBackend(Backend[torch.Tensor]):
         :param perm: e.g. [0, 2, 1]
         :return: permuted (transposed) raw tensor; wraps torch.permute
         """
+        if all(p == i for i, p in enumerate(perm)):
+            return raw_tensor
         return torch.permute(raw_tensor, tuple(perm))
 
     @staticmethod
@@ -786,6 +788,26 @@ class TorchBackend(Backend[torch.Tensor]):
             out.raw_tensor = torch.narrow(source.raw_tensor, dim=axis_int, start=start, length=end - start)
         return out
 
+    @staticmethod
+    def where(
+        cond: Tensor,
+        true_: Union[Tensor, rf.RawTensorTypes],
+        false_: Union[Tensor, rf.RawTensorTypes],
+        *,
+        allow_broadcast_all_sources: bool = False,
+    ) -> Tensor:
+        """where"""
+        true_ = rf.convert_to_tensor(true_, _backend=TorchBackend)
+        false_ = rf.convert_to_tensor(false_, _backend=TorchBackend)
+        out = Tensor.get_common_data([true_, false_, cond], allow_broadcast_all_sources=allow_broadcast_all_sources)
+        out.dtype = true_.dtype
+        out.sparse_dim = true_.sparse_dim
+        cond_bc = cond.copy_compatible_to(out, check_dtype=False, check_sparse=False)
+        true_bc = true_.copy_compatible_to(out, check_dtype=False, check_sparse=False)
+        false_bc = false_.copy_compatible_to(out, check_dtype=False, check_sparse=False)
+        out.raw_tensor = torch.where(cond_bc.raw_tensor, true_bc.raw_tensor, false_bc.raw_tensor)
+        return out
+
     @staticmethod
     def matmul(a: _TT, b: _TT, *, reduce: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> _TT:
         """
@@ -1265,20 +1287,25 @@ class TorchBackend(Backend[torch.Tensor]):
         batch_dims = [d for d in source.dims if d not in (in_dim,) + tuple(in_spatial_dims)]
         # Torch conv expects (N,C,<spatial dims>) as shape.
         source = source.copy_transpose(batch_dims + [in_dim] + list(in_spatial_dims))
-        src_raw = torch.reshape(
-            source.raw_tensor,
-            # potentially merge batch dims all together
-            [-1, in_dim.get_dim_value()] + [d.get_dim_value() for d in in_spatial_dims],
-        )
+        if len(batch_dims) == 1:
+            src_raw = source.raw_tensor
+        else:
+            src_raw = torch.reshape(
+                source.raw_tensor,
+                # potentially merge batch dims all together
+                [-1, in_dim.get_dim_value()] + [d.get_dim_value() for d in in_spatial_dims],
+            )
         use_striding = strides and (strides > 1 if isinstance(strides, int) else any(s > 1 for s in strides))
-        if (padding == "same" and use_striding) or torch.onnx.is_in_onnx_export():
+        if padding == "same" and not use_striding and all(d.dimension % 2 == 1 for d in filter_size):
+            if all(filter_size[0].dimension == d.dimension for d in filter_size):  # all same
+                padding = (filter_size[0].dimension - 1) // 2
+            else:
+                padding = tuple(d.dimension // 2 for d in filter_size)
+        if padding == "same" and (use_striding or torch.onnx.is_in_onnx_export()):
             # padding='same' is not supported for strided convolutions.
             # Moreover, padding specified as a string isn't supported for ONNX exporting as of 2023/05/19.
-            if torch.onnx.is_in_onnx_export():
-                # Explicit padding will be calculated below.
-                padding = 0
-            else:
-                padding = "valid"
+            # Manual add the padding, and then do not use any padding in the conv.
+            padding = 0
             pads = []
             for i, s in reversed(list(enumerate(filter_size))):
                 if use_striding:
@@ -1313,6 +1340,9 @@ class TorchBackend(Backend[torch.Tensor]):
                 pad_right = pad - pad_left
                 pads.extend([pad_left, pad_right])
             src_raw = torch.nn.functional.pad(src_raw, pads)
+        if padding == "valid":
+            # padding as string is not supported e.g. in ONNX.
+            padding = 0
         if len(filter_size) == 1:
             # There is also conv_tbc, but it's a bit limited (no dilation)
             # and also unclear when exactly it is faster.
@@ -1348,7 +1378,10 @@ class TorchBackend(Backend[torch.Tensor]):
         else:
             raise ValueError(f"invalid number of filter dims {filter_size}, expected 1, 2, or 3")
         out = Tensor("conv", dims=batch_dims + [out_dim] + list(out_spatial_dims), dtype=source.dtype)
-        out.raw_tensor = torch.reshape(out_raw, [d.get_dim_value() for d in out.dims])
+        if len(batch_dims) == 1:
+            out.raw_tensor = out_raw
+        else:
+            out.raw_tensor = torch.reshape(out_raw, [d.get_dim_value() for d in out.dims])
         out.feature_dim = out_dim
         return out, out_spatial_dims
 

{returnn-1.20230519.174327 → returnn-1.20230521.10102/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/tests/test_rf_array.py

@@ -224,3 +224,22 @@ def test_shift_right():
         out.mark_as_default_output(shape=(batch_dim, time_dim))
 
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
+
+
+def test_where():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "cond": Tensor("cond", [batch_dim, time_dim], dtype="bool"),
+            "true": Tensor("true", [batch_dim, time_dim, in_dim], dtype="float32"),
+            "false": Tensor("false", [batch_dim, in_dim], dtype="float32"),
+        }
+    )
+
+    # noinspection PyShadowingNames,PyUnusedLocal
+    def _forward_step(*, model: rf.Module, extern_data: TensorDict):
+        out = rf.where(extern_data["cond"], extern_data["true"], extern_data["false"])
+        out.mark_as_default_output(shape=(batch_dim, time_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)

{returnn-1.20230519.174327 → returnn-1.20230521.10102}/tools/export_to_onnx.py

@@ -174,14 +174,7 @@ def main():
         model_outputs_raw_keys.append(k)
         for i, dim in enumerate(v.dims):
             if dim.is_batch_dim() or dim.is_dynamic():
-                # For PyTorch we usually don't have the output dimensions as Dim
-                # classes available in the final tensor, so we can't call
-                # mark_as_default_output(shape=(batch_dim, ...)).
-                # These dimensions are automatically instantiated to actual values when
-                # the output tensor is passed as parameter to mark_as_default_output
-                # without any shape. Therefore, we don't add these to the output.
-                if is_rf_module:
-                    model_outputs_raw_keys.append(f"{k}:size{i}")
+                model_outputs_raw_keys.append(f"{k}:size{i}")
 
     dynamic_axes = {}
     for k, v in list(extern_data.data.items()) + list(model_outputs.data.items()):

returnn-1.20230519.174327/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230519.174327'
-long_version = '1.20230519.174327+git.4785332'