returnn 1.20230817.13718__tar.gz → 1.20230820.31622__tar.gz

{returnn-1.20230817.13718/returnn.egg-info → returnn-1.20230820.31622}/PKG-INFO

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

returnn-1.20230820.31622/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230820.031622'
+long_version = '1.20230820.031622+git.473c400'

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn/frontend/__init__.py

@@ -29,6 +29,7 @@ from .dims import *
 from .dropout import *
 from .dtype import *
 from .gradient import *
+from .graph import *
 from .label_smoothing import *
 from .linear import *
 from .loop import *

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn/frontend/_backend.py

@@ -39,6 +39,30 @@ class Backend(Generic[T]):
         """
         raise NotImplementedError
 
+    @staticmethod
+    def get_tensor_dependencies(x: Tensor) -> Sequence[Tensor]:
+        """
+        :param x: tensor
+        :return: list of all tensors which are inputs to `x`, ancestor tensors, dependencies.
+            E.g. :func:`tf.Tensor.op.inputs`.
+            This mostly makes sense for graph-based frameworks
+            but eager-based frameworks might have this too with enabled gradient tape,
+            as they should know the inputs.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def get_tensor_consumers(x: Tensor) -> Sequence[Tensor]:
+        """
+        :param x: tensor
+        :return: list of all tensors depending on `x`, descendant tensors, used by.
+            E.g. :func:`tf.Tensor.consumers`.
+            This mostly makes sense for graph-based frameworks
+            but eager-based frameworks might have this too with enabled gradient tape,
+            as they should know the consumers.
+        """
+        raise NotImplementedError
+
     @staticmethod
     def cond(pred: Tensor, true_fn: Callable, false_fn: Callable):
         """

returnn-1.20230820.31622/returnn/frontend/graph.py

@@ -0,0 +1,65 @@
+"""
+Graph utils.
+
+They might be just available for graph-based frameworks,
+but not necessarily.
+"""
+
+from __future__ import annotations
+from typing import Union, Sequence, Callable, List
+from returnn.tensor import Tensor
+
+
+__all__ = [
+    "get_tensor_dependencies",
+    "get_tensor_consumers",
+    "walk_tensor_consumers",
+]
+
+
+def get_tensor_dependencies(x: Tensor) -> Sequence[Tensor]:
+    """
+    :param x: tensor
+    :return: list of tensors which are inputs to x
+    """
+    # noinspection PyProtectedMember
+    return x._raw_backend.get_tensor_dependencies(x)
+
+
+def get_tensor_consumers(x: Tensor) -> Sequence[Tensor]:
+    """
+    :param x: tensor
+    :return: list of tensors which consume x
+    """
+    # noinspection PyProtectedMember
+    return x._raw_backend.get_tensor_consumers(x)
+
+
+def walk_tensor_consumers(
+    seed: Union[Tensor, Sequence[Tensor]],
+    *,
+    filter_outputs: Callable[[Tensor], bool] = None,
+    ending_condition: Callable[[Tensor], bool],
+) -> List[Tensor]:
+    """
+    :param seed: tensor
+    :param filter_outputs: if given, this function will be called with each tensor,
+        and if it returns False, the tensor will be skipped.
+    :param ending_condition: if given, this function will be called with each tensor,
+        and if it returns True, we return.
+    :return: yields all tensors which consume seed, and their consumers, recursively
+    """
+    seen = set()
+    queue = [seed] if isinstance(seed, Tensor) else list(seed)
+    outputs = []
+    while queue:
+        x = queue.pop(0)
+        if x in seen:
+            continue
+        seen.add(x)
+        if not filter_outputs or filter_outputs(x):
+            outputs.append(x)
+        if ending_condition and ending_condition(x):
+            break
+        queue.extend(get_tensor_consumers(x))
+    return outputs

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn/tensor/control_flow_ctx.py

@@ -31,14 +31,16 @@ class ControlFlowContext:
         """
 
         Loop = "loop"
-        Cond = "cond"
+        CondTrue = "cond-true"
+        CondFalse = "cond-false"
 
-    def __init__(self, kind: str, outer_ctx: Optional[ControlFlowContext] = None):
+    def __init__(self, *, kind: str, identifier: str, outer_ctx: Optional[ControlFlowContext] = None):
         """
         :param kind: from ControlFlowContext.Types
         :param outer_ctx:
         """
         self.kind = kind
+        self.identifier = identifier
         self._outer_ctx = outer_ctx
         self._tf_control_flow_ctx = None  # type: Optional[_TFControlFlowCtx]
         self._loop_spatial_dim = None  # type: Optional[Dim]
@@ -142,7 +144,7 @@ class ControlFlowContext:
         """
         :rtype: bool
         """
-        return self.kind == self.Types.Cond
+        return self.kind in {self.Types.CondTrue, self.Types.CondFalse}
 
     @property
     def outer_ctx(self):

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn/tf/frontend_layers/_backend.py

@@ -3,7 +3,7 @@ High-level backend for RETURNN layers
 """
 
 from __future__ import annotations
-from typing import TypeVar, Union, Sequence, Optional, Any, Callable, Tuple, Dict
+from typing import TypeVar, Union, Sequence, Optional, Any, Callable, Tuple, Dict, List
 import contextlib
 import numpy
 import tensorflow as tf
@@ -40,6 +40,30 @@ class ReturnnLayersBackend(Backend[Layer]):
         """executing eagerly"""
         return False
 
+    @staticmethod
+    def get_tensor_dependencies(x: Tensor[Layer]) -> Sequence[Tensor]:
+        """get tensor inputs"""
+        deps: List[Tensor] = []
+        visited = set()
+        for dep in x.raw_tensor.get_tensor_dependencies():
+            if not dep.tensor or dep.tensor in visited:
+                continue
+            visited.add(dep.tensor)
+            deps.append(dep.tensor)
+        return deps
+
+    @staticmethod
+    def get_tensor_consumers(x: Tensor[Layer]) -> Sequence[Tensor]:
+        """get tensor consumers"""
+        usages: List[Tensor] = []
+        visited = set()
+        for use in x.raw_tensor.usages:
+            if not use.tensor or use.tensor in visited:
+                continue
+            visited.add(use.tensor)
+            usages.append(use.tensor)
+        return usages
+
     @staticmethod
     def cond(pred: Tensor, true_fn: Callable, false_fn: Callable):
         """cond"""

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn/tf/frontend_layers/cond.py

@@ -59,7 +59,9 @@ class Cond(Generic[T]):
         )
         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()
+            kind=ControlFlowContext.Types.CondTrue,
+            outer_ctx=self.name_ctx.control_flow_ctx(),
+            identifier=self.name_ctx.get_abs_name(),
         )
         self.true_branch_name_ctx = rfl.Layer(
             module=self.layer_module,
@@ -72,7 +74,9 @@ class Cond(Generic[T]):
         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()
+            kind=ControlFlowContext.Types.CondFalse,
+            outer_ctx=self.name_ctx.control_flow_ctx(),
+            identifier=self.true_branch_control_flow_ctx.identifier,
         )
         self.false_branch_name_ctx = rfl.Layer(
             module=self.layer_module,

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn/tf/frontend_layers/loop.py

@@ -79,14 +79,17 @@ class Loop:
             if value is not NotSpecified and value is not None and key not in {"self", "__class__", "name"}
         }
         self.layer_module = LoopModule(loop=self)
+        parent_name_ctx = rfl.Layer.current_ctx()
         self.control_flow_ctx = ControlFlowContext(
-            kind=ControlFlowContext.Types.Loop, outer_ctx=rfl.Layer.inner_control_flow()
+            kind=ControlFlowContext.Types.Loop,
+            outer_ctx=rfl.Layer.inner_control_flow(),
+            identifier=parent_name_ctx.get_abs_name(),
         )
         self.control_flow_ctx.loop_spatial_dim = axis
         self.name_ctx = rfl.Layer(
             module=self.layer_module,
             suggested_name=name,
-            parent=rfl.Layer.current_ctx(),
+            parent=parent_name_ctx,
             new_control_flow_ctx=self.control_flow_ctx,
             can_access_children=False,
         )

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn/tf/layers/base.py

@@ -175,6 +175,9 @@ class LayerBase(object):
         :param str|callable|None custom_param_importer: used by :func:`set_param_values_by_dict`
         :param str|None register_as_extern_data: registers output in network.extern_data
         :param None|((LayerBase)->list[tf.Operation]) control_dependencies_on_output:
+            This is mostly to perform some checks *after* the layer output has been computed,
+            *before* the layer output is used anywhere else.
+            There is also the :class:`IdentityLayer` with the option ``control_dependencies``.
         :param None|bool|dict[str] debug_print_layer_output: same as global config option but per layer
         :param str _name: just for internal construction, should be the same as ``name``
         :param returnn.tf.network.TFNetwork _network: just for internal construction, should be the same as ``network``

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn/tf/layers/basic.py

@@ -9157,7 +9157,7 @@ class CombineLayer(LayerBase):
             out_type_["dim"] = n_out
         if out_type:
             if isinstance(out_type, dict):
-                if "shape" in out_type or "dim_tags" in out_type:
+                if "shape" in out_type or "dim_tags" in out_type or "dims" in out_type:
                     out_type_.pop("dims", None)
                     out_type_.pop("batch_dim_axis", None)
                     out_type_.pop("feature_dim_axis", None)
@@ -9588,6 +9588,15 @@ class CondLayer(LayerBase):
     See also :class:`SwitchLayer`, which uses :func:`tf.where`.
     Here, we use `tf.cond` instead. I.e. the condition has to be a scalar bool,
     and only the corresponding true/false branch is computed.
+
+    ``true_layer``/``false_layer`` are layer dicts, which are in the same namescope as this layer,
+    however, they are in the corresponding control flow context (tf.cond).
+
+    You can use :class:`SubnetworkLayer` inside to embed any more complex logic.
+
+    There can be more than one output via sub-layers.
+    Specifically, it will make all from :func:`get_available_sub_layer_names` available.
+    In :class:`SubnetworkLayer`, that are all the output layers in the sub-network.
     """
 
     layer_class = "cond"

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn/tf/layers/rec.py

@@ -1360,7 +1360,9 @@ class _SubnetworkRecCell(object):
         from returnn.tf.util.data import ControlFlowContext
 
         control_flow_ctx = ControlFlowContext(
-            kind=ControlFlowContext.Types.Loop, outer_ctx=parent_net.get_control_flow_ctx()
+            kind=ControlFlowContext.Types.Loop,
+            outer_ctx=parent_net.get_control_flow_ctx(),
+            identifier="%s%s" % (parent_net.get_absolute_name_scope_prefix(), rec_layer_name),
         )
         control_flow_ctx.loop_spatial_dim = time_dim_tag
         self.net = TFNetwork(

{returnn-1.20230817.13718 → returnn-1.20230820.31622/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/returnn.egg-info/SOURCES.txt

@@ -168,6 +168,7 @@ returnn/frontend/dims.py
 returnn/frontend/dropout.py
 returnn/frontend/dtype.py
 returnn/frontend/gradient.py
+returnn/frontend/graph.py
 returnn/frontend/init.py
 returnn/frontend/label_smoothing.py
 returnn/frontend/linear.py

{returnn-1.20230817.13718 → returnn-1.20230820.31622}/tests/test_TFUtil.py

@@ -1653,7 +1653,7 @@ def test_dim_math_pad_conv_valid_in_ctx():
     time = SpatialDim("time:var:extern_data:data")
     loop_dim = SpatialDim("time:var:extern_data:classes")
     batch_info = BatchInfo.make_global_batch_info(-1)
-    ctx = ControlFlowContext(kind=ControlFlowContext.Types.Loop)
+    ctx = ControlFlowContext(kind=ControlFlowContext.Types.Loop, identifier="loop")
     ctx.loop_spatial_dim = loop_dim
     time_ = time.get_for_batch_ctx(batch=batch_info, ctx=ctx)
     # Note: once time is actually defined (dyn_size_ext is set), the following assert would not be the case,
@@ -1676,7 +1676,7 @@ def test_dim_math_pad_conv_valid_in_ctx_derived():
 
     loop_dim = SpatialDim("time:var:extern_data:classes")
     batch_info = BatchInfo.make_global_batch_info(-1)
-    ctx = ControlFlowContext(kind=ControlFlowContext.Types.Loop)
+    ctx = ControlFlowContext(kind=ControlFlowContext.Types.Loop, identifier="loop")
     ctx.loop_spatial_dim = loop_dim
     time_undefined = Dim(kind=Dim.Types.Spatial, description="time_undefined", dimension=None)
     time_ = time_undefined.get_for_batch_ctx(batch=batch_info, ctx=ctx)

returnn-1.20230817.13718/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230817.013718'
-long_version = '1.20230817.013718+git.576322b'