returnn 1.20230818.131937__tar.gz → 1.20230820.31622__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230818.131937
+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.20230818.131937 → 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.20230818.131937 → 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.20230818.131937 → 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.20230818.131937 → 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.20230818.131937 → 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.20230818.131937 → returnn-1.20230820.31622/returnn.egg-info}/PKG-INFO

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

@@ -1,2 +0,0 @@
-version = '1.20230818.131937'
-long_version = '1.20230818.131937+git.059fc8a'