returnn 1.20230519.215329__tar.gz → 1.20230521.10102__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230519.215329
+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.215329 → 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.215329 → 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.215329 → 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.215329 → 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.215329 → 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.215329 → returnn-1.20230521.10102}/returnn/torch/frontend/_backend.py

@@ -788,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:
         """

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230519.215329
+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.215329 → 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.215329/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230519.215329'
-long_version = '1.20230519.215329+git.927c866'