returnn 1.20230816.2116__tar.gz → 1.20230816.20915__tar.gz

{returnn-1.20230816.2116/returnn.egg-info → returnn-1.20230816.20915}/PKG-INFO

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

returnn-1.20230816.20915/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230816.020915'
+long_version = '1.20230816.020915+git.37ebe2e'

{returnn-1.20230816.2116 → returnn-1.20230816.20915}/returnn/frontend/array_.py

@@ -73,12 +73,12 @@ def convert_to_tensor(
         if dims is None:
             dims = ()
         if dtype is None:
-            if isinstance(value, int):
+            if isinstance(value, bool):
+                dtype = "bool"
+            elif isinstance(value, int):
                 dtype = rf.get_default_int_dtype()
             elif isinstance(value, float):
                 dtype = rf.get_default_float_dtype()
-            elif isinstance(value, bool):
-                dtype = "bool"
             elif isinstance(value, str):
                 dtype = "string"
             elif isinstance(value, numpy.number):
@@ -527,7 +527,7 @@ def reverse_sequence(tensor: Tensor, *, axis: Dim) -> Tensor:
 
 
 def where(
-    cond: Tensor,
+    cond: Union[Tensor, rf.RawTensorTypes],
     true_: Union[Tensor, rf.RawTensorTypes],
     false_: Union[Tensor, rf.RawTensorTypes],
     *,
@@ -538,6 +538,7 @@ def where(
 
     :return: true_ if cond else false_, elemwise.
     """
+    cond = rf.convert_to_tensor(cond)
     # noinspection PyProtectedMember
     return cond._raw_backend.where(cond, true_, false_, allow_broadcast_all_sources=allow_broadcast_all_sources)
 

{returnn-1.20230816.2116 → returnn-1.20230816.20915}/returnn/frontend/loop.py

@@ -146,7 +146,7 @@ def scan(
             seq_len_ = seq_len_ + rf.cast(prev_cond, dtype=seq_len_.dtype)
             y, s = body(None, s)
             tree.assert_same_structure(ys_, y)
-            ys_ = tree.map_structure(lambda ys__, y_: ys__.push_back(y_), ys_, y)
+            ys_ = tree.map_structure(lambda ys__, y_: ys__.push_back(y_) if ys__ is not None else None, ys_, y)
             c = cond(None, s)
             c = rf.logical_and(c, prev_cond)
             return i + 1, seq_len_, c, s, ys_
@@ -168,7 +168,7 @@ def scan(
                 rf.constant(0, dtype=rf.get_default_array_index_dtype(), dims=cond_dims),  # seq_len
                 initial_cond,  # initial cond. keep this in state such that we can update seq_len in body
                 initial,  # state
-                tree.map_structure(lambda y: TensorArray(y), ys),
+                tree.map_structure(lambda y: TensorArray(y) if y is not None else None, ys),
             ),
         )
 
@@ -188,7 +188,7 @@ def scan(
             i, s, ys_ = _s
             y, s = body(tree.map_structure(lambda x: x[i], xs), s)
             tree.assert_same_structure(ys_, y)
-            ys_ = tree.map_structure(lambda ys__, y_: ys__.push_back(y_), ys_, y)
+            ys_ = tree.map_structure(lambda ys__, y_: ys__.push_back(y_) if ys__ is not None else None, ys_, y)
             return i + 1, s, ys_
 
         _, final_s, ys = while_loop(
@@ -197,12 +197,12 @@ def scan(
             (
                 rf.constant(0, dtype=rf.get_default_array_index_dtype(), dims=()),  # i
                 initial,  # state
-                tree.map_structure(lambda y: TensorArray(y), ys),
+                tree.map_structure(lambda y: TensorArray(y) if y is not None else None, ys),
             ),
         )
 
     if not return_tensor_arrays:
-        ys = tree.map_structure(lambda ys_: ys_.stack(axis=spatial_dim), ys)
+        ys = tree.map_structure(lambda ys_: ys_.stack(axis=spatial_dim) if ys_ is not None else None, ys)
     return ys, final_s, spatial_dim
 
 

{returnn-1.20230816.2116 → returnn-1.20230816.20915}/returnn/frontend/rand.py

@@ -50,6 +50,7 @@ from __future__ import annotations
 from typing import Optional, Union, Sequence, Dict
 import numpy
 from returnn.tensor import Tensor, Dim
+import returnn.frontend as rf
 from ._backend import global_backend as _global_backend
 
 
@@ -201,6 +202,8 @@ def random(
         if static:
             if seed is None:
                 seed = get_static_step_based_seed()
+    if dtype is None:
+        dtype = rf.get_default_float_dtype()
     return _global_backend.random(
         dims=dims,
         dtype=dtype,

{returnn-1.20230816.2116 → returnn-1.20230816.20915}/returnn/frontend/tensor_array.py

@@ -37,6 +37,7 @@ class TensorArray:
         _backend: Optional[Type[Backend]] = None,
         _enable_delayed_check: bool = False,
     ):
+        assert isinstance(tensor_template, Tensor)
         self.tensor_template = tensor_template
         if _backend is None:
             _backend = global_backend

{returnn-1.20230816.2116 → returnn-1.20230816.20915}/returnn/tf/frontend_layers/_backend.py

@@ -688,7 +688,6 @@ class ReturnnLayersBackend(Backend[Layer]):
         indices = []
         for i, a in enumerate(axes):
             assert isinstance(a, Dim)
-            assert not a.need_masking()  # not supported currently
             sub_name = "indices" if single_axis else f"indices{i}"
             indices_data = values.copy_template(name=f"{values.name}_{sub_name}_{a.description}")
             indices_data.dtype = "int32"

{returnn-1.20230816.2116 → returnn-1.20230816.20915}/returnn/tf/layers/basic.py

@@ -9850,7 +9850,7 @@ class TopKLayer(LayerBase):
         :param bool sorted:
         """
         super(TopKLayer, self).__init__(**kwargs)
-        in_data = self.sources[0].output
+        in_data = self.sources[0].output.copy()
 
         if isinstance(axis, (str, Dim)):
             single_axis = True
@@ -9865,6 +9865,14 @@ class TopKLayer(LayerBase):
         in_data = in_data.copy_transpose([in_data.get_axis_from_description(a) for a in remaining_axes + axes])
         assert in_data.dim_tags == tuple(remaining_axes + axes)
 
+        for a in axes:
+            if a.need_masking():
+                in_data.raw_tensor = tf_util.where_bc(
+                    a.get_mask(dim_order=in_data.dims, device=in_data.device).raw_tensor,
+                    in_data.raw_tensor,
+                    in_data.raw_tensor.dtype.min,
+                )
+
         # Merge the axes because top_k will do a joint search over them.
         if len(axes) > 1:
             merged_axis = axes[0]

{returnn-1.20230816.2116 → returnn-1.20230816.20915}/returnn/torch/frontend/_backend.py

@@ -1043,7 +1043,7 @@ class TorchBackend(Backend[torch.Tensor]):
     # noinspection PyShadowingBuiltins
     @staticmethod
     def top_k(
-        source: Tensor,
+        source: Tensor[torch.Tensor],
         *,
         axis: Union[Dim, Sequence[Dim]],
         k: Union[int, Tensor],
@@ -1053,6 +1053,21 @@ class TorchBackend(Backend[torch.Tensor]):
         """top_k"""
         if not k_dim:
             k_dim = Dim(k, name="top-k-dim")
+        axes = [axis] if isinstance(axis, Dim) else axis
+        if any(a.need_masking() for a in axes):
+            mask_value = (
+                torch.finfo(source.raw_tensor.dtype).min
+                if source.raw_tensor.dtype.is_floating_point
+                else torch.iinfo(source.raw_tensor.dtype).min
+            )
+            source = source.copy()
+            for a in axes:
+                if a.need_masking():
+                    source = rf.where(
+                        a.get_mask(dim_order=source.dims, device=source.device),
+                        source,
+                        mask_value,
+                    )
         if isinstance(axis, (list, tuple)):
             # Move axis to the end, in the right order.
             source = source.copy_transpose([d for d in source.dims if d not in axis] + list(axis))
@@ -1078,7 +1093,6 @@ class TorchBackend(Backend[torch.Tensor]):
                 indices_out.insert(0, indices)
             return values, indices_out, k_dim
         assert isinstance(axis, Dim)
-        assert not axis.need_masking()  # not supported currently
         axis_int = source.get_axis_from_description(axis, allow_int=False)
         values_raw, indices_raw = torch.topk(
             source.raw_tensor, k=k_dim.get_dim_value(), dim=axis_int, largest=True, sorted=sorted

{returnn-1.20230816.2116 → returnn-1.20230816.20915/returnn.egg-info}/PKG-INFO

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

returnn-1.20230816.2116/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230816.002116'
-long_version = '1.20230816.002116+git.e1ca657'