returnn 1.20230501.124840__tar.gz → 1.20230503.142906__tar.gz

{returnn-1.20230501.124840/returnn.egg-info → returnn-1.20230503.142906}/PKG-INFO

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

returnn-1.20230503.142906/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230503.142906'
+long_version = '1.20230503.142906+git.1c628d2'

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/frontend/_backend.py

@@ -3,7 +3,7 @@ Backends for the frontend API
 """
 
 from __future__ import annotations
-from typing import TYPE_CHECKING, Optional, Any, Union, TypeVar, Generic, Type, Sequence, Dict, Tuple
+from typing import TYPE_CHECKING, Optional, Any, Union, TypeVar, Generic, Type, Callable, Sequence, Dict, Tuple
 import contextlib
 import numpy
 import returnn.frontend as rf
@@ -40,6 +40,19 @@ class Backend(Generic[T]):
         """
         raise NotImplementedError
 
+    @staticmethod
+    def cond(pred: Tensor, true_fn: Callable, false_fn: Callable):
+        """
+        cond: conditional execution.
+
+        Note that this does not need an implementation for eager-based frameworks
+        (:func:`executing_eagerly` returns True),
+        as the :func:`returnn.frontend.cond` function already covers that case.
+        """
+        # noinspection PyProtectedMember
+        assert not pred._raw_backend.executing_eagerly(), "should not get here"
+        raise NotImplementedError
+
     @staticmethod
     def set_random_seed(seed: int):
         """

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/frontend/cond.py

@@ -8,11 +8,14 @@ https://github.com/rwth-i6/returnn/issues/1282
 """
 
 from __future__ import annotations
-from typing import Union
+from typing import Union, TypeVar, Callable
 from returnn.tensor import Tensor
 
 
-def cond(pred: Union[bool, Tensor], true_fn, false_fn):
+T = TypeVar("T")
+
+
+def cond(pred: Union[bool, Tensor], true_fn: Callable[[], T], false_fn: Callable[[], T]) -> T:
     """
     :param pred:
     :param true_fn:
@@ -24,6 +27,7 @@ def cond(pred: Union[bool, Tensor], true_fn, false_fn):
             return true_fn()
         else:
             return false_fn()
+    assert isinstance(pred, Tensor) and pred.dims == () and pred.dtype == "bool"
     # noinspection PyProtectedMember
     backend = pred._raw_backend
     if backend.executing_eagerly():
@@ -31,4 +35,4 @@ def cond(pred: Union[bool, Tensor], true_fn, false_fn):
             return true_fn()
         else:
             return false_fn()
-    raise NotImplementedError("cond() not implemented for backend %r" % backend)
+    return backend.cond(pred, true_fn, false_fn)

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/frontend/const.py

@@ -10,7 +10,7 @@ from ._backend import global_backend
 import returnn.frontend as rf
 
 
-__all__ = ["full", "constant", "fill", "zeros", "ones"]
+__all__ = ["full", "constant", "fill", "zeros", "ones", "zeros_like", "ones_like"]
 
 
 def full(
@@ -54,3 +54,13 @@ def ones(dims: Sequence[Dim], *, dtype: Optional[str] = None, sparse_dim: Option
     ones. float by default.
     """
     return full(dims=dims, fill_value=1, dtype=dtype or rf.get_default_float_dtype(), sparse_dim=sparse_dim)
+
+
+def zeros_like(other: Tensor) -> Tensor:
+    """zeros like other"""
+    return zeros(dims=other.dims, dtype=other.dtype, sparse_dim=other.sparse_dim)
+
+
+def ones_like(other: Tensor) -> Tensor:
+    """ones like other"""
+    return ones(dims=other.dims, dtype=other.dtype, sparse_dim=other.sparse_dim)

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/frontend/parameter.py

@@ -138,7 +138,12 @@ class Parameter(Tensor[T]):
     def trainable(self, trainable: Optional[bool]):
         self._trainable = trainable
         if trainable is None:
-            trainable = not self.auxiliary
+            if self.auxiliary:
+                trainable = False
+            elif self.dtype.startswith("int"):
+                trainable = False
+            else:
+                trainable = True
         self._raw_backend.set_parameter_trainable(self, trainable)
 
     @property

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/tf/frontend_layers/__init__.py

@@ -5,6 +5,7 @@ TensorFlow frontend.
 from ._backend import ReturnnLayersBackend
 
 # Sorted by name.
+from .cond import Cond
 from .config_entry_points import *
 from .debug_eager_mode import *
 from .dims import *

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/tf/frontend_layers/_backend.py

@@ -3,7 +3,7 @@ High-level backend for RETURNN layers
 """
 
 from __future__ import annotations
-from typing import Union, Sequence, Optional, Any, Tuple, Dict
+from typing import Union, Sequence, Optional, Any, Callable, Tuple, Dict
 import contextlib
 import numpy
 import tensorflow as tf
@@ -37,6 +37,14 @@ class ReturnnLayersBackend(Backend[Layer]):
         """executing eagerly"""
         return False
 
+    @staticmethod
+    def cond(pred: Tensor, true_fn: Callable, false_fn: Callable):
+        """cond"""
+        with rfl.Cond(pred) as cond:
+            cond.true = true_fn()
+            cond.false = false_fn()
+        return cond.result
+
     @staticmethod
     def set_random_seed(seed: int):
         """

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/tf/frontend_layers/_utils.py

@@ -30,6 +30,25 @@ def copy(tensor: Tensor[rfl.Layer], *, name: Union[rfl.Layer, str]) -> Tensor[rf
     return rfl.make_layer({"class": "copy", "from": tensor}, name=name)
 
 
+def mark_as_output_in_scope(tensor: Tensor, scope: rfl.Layer) -> Tensor:
+    """
+    Mark this as an output.
+    """
+    assert tensor.raw_tensor.layer_dict, f"mark_as_output can only be called on a layer, not a layer-ref {tensor}."
+    res = tensor
+    if tensor.raw_tensor is scope.children.get("output"):
+        pass  # not needed
+    elif tensor.raw_tensor.parent is not scope:
+        res = copy(tensor, name=scope.get_new_child(suggested_name=tensor.raw_tensor.get_abs_name(join_str="_")))
+        res.raw_tensor.layer_dict["is_output_layer"] = True
+    else:
+        assert tensor.raw_tensor.parent is scope
+        assert tensor.raw_tensor.layer_dict
+        tensor.raw_tensor.layer_dict["is_output_layer"] = True
+    scope.marked_outputs.append(res)
+    return res
+
+
 def get_last_hidden_state(
     source: Tensor,
     *,

returnn-1.20230503.142906/returnn/tf/frontend_layers/cond.py

@@ -0,0 +1,252 @@
+"""
+Conditional logic
+
+https://github.com/rwth-i6/returnn_common/issues/24
+"""
+
+from __future__ import annotations
+from typing import List, TypeVar, Generic
+from tensorflow.python.util import nest
+from returnn.tensor import Tensor, ControlFlowContext
+import returnn.frontend as rf
+import returnn.tf.frontend_layers as rfl
+from . import _utils
+
+
+T = TypeVar("T")
+
+
+__all__ = ["Cond", "CondModule"]
+
+
+class Cond(Generic[T]):
+    """
+    Conditional branching. Basically behaves like ``if ... else ...``.
+    Only one branch will be executed, and the condition needs to be a bool scalar.
+    This wraps to :class:`CondLayer` in RETURNN and to ``tf.cond`` in TensorFlow.
+
+    Example::
+
+        with Cond(cond) as cond_obj:
+          cond_obj.true = mod_true_case(x)
+          cond_obj.false = mod_false_case(x)
+          y = cond_obj.result
+
+    Corresponds to::
+
+        if cond:
+          y = mod_true_case(x)
+        else:
+          y = mod_false_case(x)
+
+    The context scope has two states corresponding to the True and False computation branch.
+    The initial state is the True branch.
+    Assigning ``cond_obj.true`` has the side effect of switching the computation to the False branch.
+    """
+
+    def __init__(self, condition: Tensor, *, name: str = "cond"):
+        self.condition = condition
+        self._entered = False
+        self._entered_state = True
+        self._true_value = None
+        self._false_value = None
+        self._result_value = None
+        self.layer_module = CondModule(cond=self)
+        self.name_ctx = rfl.Layer(
+            module=self.layer_module, suggested_name=name, parent=rfl.Layer.current_ctx(), can_access_children=False
+        )
+        self.name_ctx.custom_layer_name_scope = ""
+        self.true_branch_control_flow_ctx = ControlFlowContext(
+            kind=ControlFlowContext.Types.Cond, outer_ctx=self.name_ctx.control_flow_ctx()
+        )
+        self.true_branch_name_ctx = rfl.Layer(
+            module=self.layer_module,
+            suggested_name="true",
+            parent=self.name_ctx,
+            virtual=True,
+            can_access_children=False,
+            new_control_flow_ctx=self.true_branch_control_flow_ctx,
+        )
+        self.true_branch_name_ctx.is_subnet = True
+        self.true_branch_name_ctx.extend_reserved_names({"output"})
+        self.false_branch_control_flow_ctx = ControlFlowContext(
+            kind=ControlFlowContext.Types.Cond, outer_ctx=self.name_ctx.control_flow_ctx()
+        )
+        self.false_branch_name_ctx = rfl.Layer(
+            module=self.layer_module,
+            suggested_name="false",
+            parent=self.name_ctx,
+            virtual=True,
+            can_access_children=False,
+            new_control_flow_ctx=self.false_branch_control_flow_ctx,
+        )
+        self.false_branch_name_ctx.is_subnet = True
+        self.false_branch_name_ctx.extend_reserved_names({"output"})
+
+    def __repr__(self):
+        return f"Cond{self.name_ctx}"
+
+    def __enter__(self):
+        assert not self._entered, f"{self} cannot enter twice"
+        self._entered = True
+        self._entered_state = True
+        self.true_branch_name_ctx.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # First exit any scopes and do cleanup without throwing any exceptions.
+        if self._entered:
+            if self._true_value is None:
+                self.true_branch_name_ctx.__exit__(exc_type, exc_val, exc_tb)
+            elif self._false_value is None:
+                self.false_branch_name_ctx.__exit__(exc_type, exc_val, exc_tb)
+        if not exc_type:  # only do error checking if there was no other exception
+            assert self._entered
+            assert self._true_value is not None, f"{self} you need to call else_()"
+            assert self._false_value is not None, f"{self} you need to call end()"
+        self._entered = False
+
+    @property
+    def true(self) -> T:
+        """
+        The getter usually would not be used.
+        """
+        return self._true_value
+
+    @true.setter
+    def true(self, true_value: T):
+        """
+        Defines the True branch value.
+        Enter the False branch.
+        Assign self.false afterwards.
+        """
+        assert self._entered, f"{self} you need to be in the context scope"
+        assert self._entered_state is True, f"{self} you cannot enter the else branch twice"
+        assert true_value is not None
+        assert self._true_value is None
+        if isinstance(true_value, Tensor):
+            true_value = _utils.copy(true_value, name=self.true_branch_name_ctx.get_child("output"))
+        else:
+            values_flat = nest.flatten(true_value)  # type: List[Tensor]
+            assert values_flat
+            for i, v in enumerate(values_flat):
+                assert isinstance(v, Tensor), f"unexpected {true_value!r}, only expects tensors, got {type(v)}"
+                if i == 0:
+                    values_flat[i] = _utils.copy(v, name=self.true_branch_name_ctx.get_child("output"))
+                else:
+                    values_flat[i] = _utils.mark_as_output_in_scope(v, scope=self.true_branch_name_ctx)
+            true_value = nest.pack_sequence_as(true_value, values_flat)
+        self.true_branch_name_ctx.__exit__(None, None, None)
+        self.false_branch_name_ctx.__enter__()
+        self._true_value = true_value
+        self._entered_state = False
+
+    @property
+    def false(self) -> T:
+        """
+        The getter usually would not be used.
+        """
+        return self._false_value
+
+    @false.setter
+    def false(self, false_value: T):
+        """
+        Define the False branch value.
+        After this, self.result is available.
+        """
+        assert self._entered, f"{self} you need to be in the context scope"
+        assert (
+            self._entered_state is False
+        ), f"{self} you need to be in the False branch, have assigned :func:`true` before"
+        assert false_value is not None
+        assert self._false_value is None
+        nest.assert_same_structure(self._true_value, false_value)
+        # This needs to match the true() setter logic.
+        if isinstance(false_value, Tensor):
+            false_value = _utils.copy(false_value, name=self.false_branch_name_ctx.get_child("output"))
+        else:
+            true_values_flat = nest.flatten(self._true_value)  # type: List[Tensor]
+            false_values_flat = nest.flatten(false_value)  # type: List[Tensor]
+            assert false_values_flat and len(false_values_flat) == len(true_values_flat)
+            for i, (true_v, false_v) in enumerate(zip(true_values_flat, false_values_flat)):
+                assert isinstance(true_v, Tensor)
+                assert isinstance(
+                    false_v, Tensor
+                ), f"unexpected {false_value!r}, only expects tensors, got {type(false_v)}"
+                assert true_v.raw_tensor.parent is self.true_branch_name_ctx
+                name = true_v.raw_tensor.name
+                false_values_flat[i] = _utils.copy(false_v, name=self.false_branch_name_ctx.get_child(name))
+                if name != "output":
+                    false_values_flat[i].raw_tensor.layer_dict["is_output_layer"] = True
+            false_value = nest.pack_sequence_as(false_value, false_values_flat)
+        self.false_branch_name_ctx.__exit__(None, None, None)
+        self._false_value = false_value
+        self._result_value = self.layer_module()
+
+    @property
+    def result(self) -> T:
+        """
+        :return: the result, after you assigned :func:`true` and :func:`false`.
+        """
+        assert self._true_value is not None, f"{self} you need to have defined the true value"
+        assert self._false_value is not None, f"{self} you need to have defined the false value"
+        assert self._result_value is not None
+        return self._result_value
+
+
+class CondModule(rf.Module):
+    """
+    This module is used internally by :class:`Cond` to create the RETURNN :class:`CondLayer` for the conditional code.
+    This module would not be directly used by the user.
+    """
+
+    def __init__(self, cond: Cond):
+        super(CondModule, self).__init__()
+        self.cond = cond
+
+    def __call__(self):
+        """
+        Makes layer dict for this loop, i.e. a RecLayer.
+
+        :return: structure like true_value/false_value
+        """
+        name_ctx = self.cond.name_ctx
+        # noinspection PyProtectedMember
+        true_value, false_value = self.cond._true_value, self.cond._false_value
+        true_values_flat = nest.flatten(true_value)  # type: List[Tensor]
+        false_values_flat = nest.flatten(false_value)  # type: List[Tensor]
+        assert len(true_values_flat) == len(false_values_flat)
+        res = rfl.make_layer(
+            {
+                "class": "cond",
+                "from": [],
+                "condition": self.cond.condition,
+                "true_layer": {
+                    "class": "subnetwork",
+                    "from": [],
+                    "subnetwork": self.cond.true_branch_name_ctx.make_net(),
+                },
+                "false_layer": {
+                    "class": "subnetwork",
+                    "from": [],
+                    "subnetwork": self.cond.false_branch_name_ctx.make_net(),
+                },
+            },
+            name=name_ctx,
+            predefined_out_data=true_values_flat[0].copy_template(),
+        )
+
+        results = []
+        for i, (true_v, false_v) in enumerate(zip(true_values_flat, false_values_flat)):
+            assert isinstance(true_v, Tensor) and isinstance(false_v, Tensor)
+            assert true_v.raw_tensor.parent is self.cond.true_branch_name_ctx
+            name = true_v.raw_tensor.name
+            if i == 0:
+                results.append(res)
+            else:
+                # noinspection PyProtectedMember
+                results.append(rfl._get_sub_layer(res, name, data=true_v.copy_template()))
+            results[-1].raw_tensor.layer_extra_dependencies.extend(
+                (self.cond.condition.raw_tensor, true_v.raw_tensor, false_v.raw_tensor)
+            )
+        return nest.pack_sequence_as(true_value, results)

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/tf/layers/basic.py

@@ -309,6 +309,7 @@ class _ConcatInputLayer(LayerBase):
 class CopyLayer(_ConcatInputLayer):
     """
     This layer does nothing, it copies its input.
+    This is not even a ``tf.identity``. It refers to the same TF tensor.
     If multiple sources are provided, they are concatenated in the feature-dim.
     """
 
@@ -323,6 +324,9 @@ class CopyLayer(_ConcatInputLayer):
           We only have this here for the :class:`CopyLayer` because the :func:`get_out_data_from_opts`
           must know about it and define the right beam.
           Also see the option ``collocate_with``, which is different in that it does *not* add a dependency.
+          Note that this will not be real TF control dependencies,
+          but it simply sets the dependency on the layer.
+          If you want to have a real TF control dependency, use :class:`IdentityLayer`.
         """
         if in_dim and out_dim:
             assert in_dim == out_dim
@@ -415,6 +419,42 @@ class CopyLayer(_ConcatInputLayer):
             d["extra_deps"] = [get_layer(src_name) for src_name in extra_deps]
 
 
+class IdentityLayer(LayerBase):
+    """
+    Wraps ``tf.identity`` with potential control dependencies.
+
+    The difference to :class:`CopyLayer` is that this creates a new TF op (``tf.identity``),
+    which allows for potential control dependencies.
+    This is the whole purpose of this layer.
+    """
+
+    layer_class = "identity"
+
+    def __init__(self, sources: List[LayerBase], control_dependencies: Sequence[LayerBase], **kwargs):
+        super().__init__(sources=sources, **kwargs)
+        assert len(sources) == 1
+        self.control_dependencies = control_dependencies
+        with tf.control_dependencies([src.output.placeholder.op for src in self.control_dependencies]):
+            self.output.placeholder = tf.identity(self.sources[0].output.placeholder)
+
+    def get_dep_layers(self) -> List[LayerBase]:
+        """deps"""
+        return super().get_dep_layers() + list(self.control_dependencies)
+
+    @classmethod
+    def get_out_data_from_opts(cls, name: str, sources: List[LayerBase], **kwargs):
+        """out"""
+        assert sources
+        return sources[0].output.copy(name="%s_output" % name)
+
+    @classmethod
+    def transform_config_dict(cls, d, network, get_layer):
+        """transform"""
+        super().transform_config_dict(d, network=network, get_layer=get_layer)
+        assert isinstance(d.get("control_dependencies"), (list, tuple))
+        d["control_dependencies"] = [get_layer(src_name) for src_name in d["control_dependencies"]]
+
+
 class ConcatLayer(LayerBase):
     """
     Concatenates the inputs in specified axes.

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/tf/util/basic.py

@@ -5850,11 +5850,13 @@ def same_control_flow_ctx(x):
         with tf.control_dependencies(None) as dep:  # this will reset the context
             yield dep
         return
-    # noinspection PyProtectedMember
-    graph._set_control_flow_context(ctx)
-    yield ctx
-    # noinspection PyProtectedMember
-    graph._set_control_flow_context(cur_ctx)
+    try:
+        # noinspection PyProtectedMember
+        graph._set_control_flow_context(ctx)
+        yield ctx
+    finally:
+        # noinspection PyProtectedMember
+        graph._set_control_flow_context(cur_ctx)
 
 
 def get_protobuf_fields(obj):

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn/torch/frontend/_backend.py

@@ -483,8 +483,12 @@ class TorchBackend(Backend[torch.Tensor]):
         :return: parameter
         """
         assert all(d.is_static() for d in tensor.dims)
-        data = torch.zeros(*(d.dimension for d in tensor.dims), dtype=TorchBackend.as_dtype_raw(tensor.dtype))
-        return torch.nn.Parameter(data)
+        data = torch.zeros([d.dimension for d in tensor.dims], dtype=TorchBackend.as_dtype_raw(tensor.dtype))
+        if tensor.dtype.startswith("int"):
+            requires_grad = False
+        else:
+            requires_grad = True
+        return torch.nn.Parameter(data, requires_grad=requires_grad)
 
     @staticmethod
     def set_parameter_initial_value(param: rf.Parameter, value: Union[None, Tensor, rf.RawTensorTypes]) -> None:
@@ -498,7 +502,8 @@ class TorchBackend(Backend[torch.Tensor]):
         assert isinstance(raw_param, torch.nn.Parameter)
         with torch.no_grad():
             if isinstance(value, Tensor):
-                raw_param.copy_(value.raw_tensor)
+                value_ = value.copy_compatible_to(param)
+                raw_param.copy_(value_.raw_tensor)
             elif isinstance(value, numpy.ndarray):
                 raw_param.copy_(torch.from_numpy(value))
             else:

{returnn-1.20230501.124840 → returnn-1.20230503.142906/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230501.124840 → returnn-1.20230503.142906}/returnn.egg-info/SOURCES.txt

@@ -220,6 +220,7 @@ returnn/tf/updater.py
 returnn/tf/frontend_layers/__init__.py
 returnn/tf/frontend_layers/_backend.py
 returnn/tf/frontend_layers/_utils.py
+returnn/tf/frontend_layers/cond.py
 returnn/tf/frontend_layers/config_entry_points.py
 returnn/tf/frontend_layers/debug_eager_mode.py
 returnn/tf/frontend_layers/dims.py
@@ -314,6 +315,7 @@ tests/test_hdf_dump.py
 tests/test_rf_array.py
 tests/test_rf_attention.py
 tests/test_rf_base.py
+tests/test_rf_cond.py
 tests/test_rf_container.py
 tests/test_rf_conv.py
 tests/test_rf_encoder_conformer.py