
"""
This module contains the layer base class :class:`LayerBase`.
"""

from __future__ import print_function

import typing
import contextlib
import tensorflow as tf
from returnn.util.basic import NotSpecified, CollectionReadCheckCovered
import returnn.tf.compat as tf_compat
import returnn.tf.util.basic as tf_util
from returnn.tf.util.data import Data, SearchBeam
from returnn.tf.util.basic import OutputWithActivation, CustomUpdate, reuse_name_scope
from returnn.log import log


class LayerBase(object):
  """
  This is the base class for all layers.
  Every layer by default has a list of source layers `sources` and defines `self.output` which is of type :class:`Data`.
  It shares some common functionality across all layers, such as explicitly defining the output format,
  some parameter regularization, and more.

  If you want to implement your own layer::

      class YourOwnLayer(_ConcatInputLayer):  # e.g. either _ConcatInputLayer or LayerBase as a base
          " some docstring "
          layer_class = "your_layer_name"

          def __init__(self, your_kwarg1, your_opt_kwarg2=None, **kwargs):
              " docstring, document the args! "
              super(YourOwnLayer, self).__init__(**kwargs)
              # Now we need to set self.output, which must be of type :class:`Data`.
              # It is set at this point to whatever we got from `self.get_out_data_from_opts()`,
              # so it is enough if we set self.output.placeholder and self.output.size_placeholder,
              # but we could also reset self.output.
              self.output.placeholder = self.input_data.placeholder + 42  # whatever you want to do
              # If you don't modify the sizes (e.g. sequence-length), just copy the input sizes.
              self.output.size_placeholder = self.input_data.size_placeholder.copy()

          @classmethod
          def get_out_data_from_opts(cls, **kwargs):
              " This is supposed to return a :class:`Data` instance as a template, given the arguments. "
              # example, just the same as the input:
              return get_concat_sources_data_template(kwargs["sources"], name="%s_output" % kwargs["name"])

  """

  layer_class = None  # type: typing.Optional[str]  # for get_layer_class()
  recurrent = False  # if the order in the time-dimension is relevant
  allow_inf_in_output = False

  # For compatibility, we have some parameter names (e.g. "L2") which do not conform to PEP8.
  # noinspection PyPep8Naming
  def __init__(self, name, network, output=None, n_out=NotSpecified, out_type=None, sources=(),
               target=None, _target_layers=None, loss=None, size_target=None,
               reuse_params=None,
               param_device=None,
               is_output_layer=None, only_on_eval=False, only_on_search=False,
               copy_output_loss_from_source_idx=None,
               batch_norm=False,
               L2=None, darc1=None,
               spatial_smoothing=0.0,
               param_variational_noise=None,
               updater_opts=None,
               initial_output=None,
               rec_previous_layer=None,
               collocate_with=None,
               trainable=True,
               custom_param_importer=None,
               register_as_extern_data=None,
               control_dependencies_on_output=None,
               _src_common_search_choices=None):
    """
    Usually the arguments, when specified in the network dict,
    are going through :func:`transform_config_dict`, before they are passed to here.
    See :func:`TFNetwork.construct_from_dict`.

    :param str name:
    :param returnn.tf.network.TFNetwork network:
    :param Data|None output: Set a specific output instead of using :func:`get_out_data_from_opts`
    :param NotSpecified|None|int n_out: output dim
    :param dict[str] out_type: kwargs for Data class. more explicit than n_out.
    :param list[LayerBase] sources: via self.transform_config_dict()
    :param str|list[str]|None target: if some loss is set, this is the target data-key,
      i.e. network.extern_data.get_data(target). alternatively, this also can be a layer name.
    :param dict[str,LayerBase]|None _target_layers: if target.startswith("layer:"), then this is target -> layer
    :param str|None size_target: like target but this is only used to set our output size in case of training
    :param Loss|None loss: via :func:`transform_config_dict`.
      Every layer can have one loss (of type :class:`Loss`), or none loss.
      In the net dict, it is specified as a string.
      In :class:`TFNetwork`, all losses from all layers will be collected.
      That is what :class:`TFUpdater.Updater` will use for training.
    :param ReuseParams|None reuse_params: if given, will opt reuse the params. see :func:`self.var_creation_scope`
    :param str|None param_device: e.g. "CPU", etc. any valid name for tf.device.
      see https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/util/device_name_utils.h
    :param float|None L2: for constraints
    :param float|None darc1: for constraints. see Generalization in Deep Learning, https://arxiv.org/abs/1710.05468
    :param float|None spatial_smoothing: see :func:`TFUtil.spatial_smoothing_energy`
    :param float|None param_variational_noise: adds variational noise to the params during training
    :param dict[str]|None updater_opts: accepts similar opts as TFUpdater, e.g. "optimizer", "learning_rate", ...
    :param bool|None is_output_layer:
    :param bool only_on_eval: if True, this layer will only be calculated in eval
    :param bool only_on_search: if True, this layer will only be calculated when search is done
    :param int|None copy_output_loss_from_source_idx: if set, will copy output_loss from this source
    :param bool|dict batch_norm: see self.batch_norm()
    :param str|float initial_output: used for recurrent layer, see self.get_rec_initial_output()
    :param LayerBase|None rec_previous_layer: via the recurrent layer, layer (template) which represents the past of us.
      You would not explicitly set this in a config. This is automatically, internally, via :class:`RecLayer`.
    :param list[str]|None collocate_with: in the rec layer, collocate with the specified other layers
    :param bool trainable: whether the parameters of this layer will be trained
    :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:
    :param None|SearchChoices _src_common_search_choices: set via :func:`SearchChoices.translate_to_common_search_beam`
    """
    self.name = name
    self.network = network
    self._register_layer()
    self.kwargs = None  # type: typing.Optional[typing.Dict[str]] # set via self.post_init
    self.target = None
    self.targets = None
    if target:
      if isinstance(target, list):
        self.targets = target
        self.target = target[0]
      else:
        assert isinstance(target, str)
        self.targets = [target]
        self.target = target
    self._target_layers = _target_layers
    self.size_target = size_target
    self.loss = loss
    if self.loss and self.loss.recurrent:
      self.recurrent = True
    if output:
      self.output = output
      if n_out is not NotSpecified:
        assert self.output.dim == n_out
      if isinstance(out_type, dict):
        if "shape" in out_type:
          assert self.output.shape == out_type["shape"]
        if "dim" in out_type:
          assert self.output.dim == out_type["dim"]
    else:
      self.output = self.get_out_data_from_opts(
        out_type=out_type, n_out=n_out,
        network=network, name=name, target=target, size_target=size_target,
        sources=sources, loss=loss)
    self.output_before_activation = None  # type: typing.Optional[OutputWithActivation]
    self.output_loss = None  # type: typing.Optional[tf.Tensor]
    if copy_output_loss_from_source_idx is not None:
      self.output_loss = sources[copy_output_loss_from_source_idx].output_loss
    self.rec_vars_outputs = {}  # type: typing.Dict[str,tf.Tensor]
    self.search_choices = None  # type: typing.Optional[SearchChoices]
    self._src_common_search_choices = _src_common_search_choices
    self._initial_output = initial_output
    self._rec_previous_layer = rec_previous_layer
    self.collocate_with = collocate_with or []
    self.post_init_hooks = []  # list of functions
    self.sources = list(sources)
    self.params = {}  # type: typing.Dict[str,tf.Variable]
    self.saveable_param_replace = {}  # type:  typing.Dict[tf.Variable,typing.Union['tensorflow.python.training.saver.BaseSaverBuilder.SaveableObject',None]]  # see get_saveable_params_dict()  # nopep8
    self.reuse_params = reuse_params
    self.param_device = param_device
    self.L2 = L2
    self.darc1 = darc1
    self.spatial_smoothing = spatial_smoothing
    self.param_variational_noise = param_variational_noise
    self.updater_opts = CollectionReadCheckCovered(updater_opts or {})
    self._is_output_layer = is_output_layer
    self.only_on_eval = only_on_eval
    self.only_on_search = only_on_search
    self.use_batch_norm = batch_norm
    self.trainable = trainable
    self.custom_param_importer = custom_param_importer
    self.control_dependencies_on_output = control_dependencies_on_output
    self.register_as_extern_data = register_as_extern_data
    # Stats will be collected by the engine.
    self.stats = {}  # type: typing.Dict[str,tf.Tensor]

  def post_init(self, layer_desc):
    """
    This gets called right after self.__init__().

    :param dict[str] layer_desc: kwargs as they are passed to self.__init__
    """
    self.kwargs = layer_desc
    if self.use_batch_norm:
      opts = {}
      if isinstance(self.use_batch_norm, dict):
        opts = self.use_batch_norm
      self.output.placeholder = self.batch_norm(self.output, **opts)
    if self.control_dependencies_on_output:
      control_deps = self.control_dependencies_on_output(self)
      if not isinstance(control_deps, (list, tuple)):
        assert isinstance(control_deps, (tf.Operation, tf.Tensor))
        control_deps = [control_deps]
      assert all([isinstance(dep, (tf.Operation, tf.Tensor)) for dep in control_deps])
      if control_deps:
        with tf.control_dependencies(control_deps):
          self.output.placeholder = tf.identity(self.output.placeholder)
    if self.register_as_extern_data:
      self.network.extern_data.extra_added_keys.add(self.register_as_extern_data)
      self.network.extern_data.data[self.register_as_extern_data] = self.output
    for func in self.post_init_hooks:
      func()

  def __repr__(self):
    return "<%s %s%r out_type=%s>" % (
      self.__class__.__name__, self.network.get_absolute_name_prefix(), self.name,
      self.output.get_description(with_name=False) if self.output else None)

  @classmethod
  def get_out_data_from_opts(cls, **kwargs):
    """
    Gets a Data template (i.e. shape etc is set but not the placeholder) for our __init__ args.
    The purpose of having this as a separate classmethod is to be able to infer the shape information
    without having to construct the layer.
    This function should not create any nodes in the computation graph.

    :param kwargs: all the same kwargs as for self.__init__()
    :return: Data template (placeholder not set)
    :rtype: Data
    """
    return cls._base_get_out_data_from_opts(**kwargs)

  @classmethod
  def _base_get_out_data_from_opts(cls, network, name, out_type=None, n_out=NotSpecified,
                                   target=None, _target_layers=None, size_target=None,
                                   sources=(), loss=None,
                                   **kwargs):
    """
    Called via BaseLayer.get_out_data_from_opts().

    :param returnn.tf.network.TFNetwork network:
    :param str name:
    :param dict[str]|None|(()->Data) out_type:
    :param int|None|NotSpecified n_out:
    :param str|list[str]|None target:
    :param dict[str,LayerBase]|None _target_layers: if target.startswith("layer:"), then this is target -> layer
    :param str|None size_target:
    :param list[LayerBase] sources:
    :param Loss|None loss:
    :param kwargs: remaining kwargs of self.__init__(), ignored here
    :return: Data template (placeholder not set)
    :rtype: Data
    """
    if callable(out_type):
      return out_type(
        network=network, name=name, n_out=n_out, target=target, size_target=size_target, sources=sources, loss=loss,
        **kwargs)
    # Any template construction should be aware of that, and eventually resolve it.
    if out_type is None:
      out_type = {}  # type: typing.Dict[str]
    else:
      out_type = out_type.copy()
    out_type.setdefault("name", "%s_output" % name)
    if "dim" not in out_type and n_out is not NotSpecified:
      out_type["dim"] = n_out
    if "dim" not in out_type and target:
      out_type["dim"] = cls._static_get_target_value(
        target=target[0] if isinstance(target, list) else target, _target_layers=_target_layers,
        network=network, mark_data_key_as_used=False).dim
    if n_out is not NotSpecified:
      assert out_type["dim"] == n_out
    sources_data = None
    if sources and sources[0]:
      sources_data = sources[0].output.copy_template()
    if sources_data and not sources_data.sparse and not out_type.get("sparse", False):
      out_type.setdefault("dtype", sources_data.dtype)
    # You are supposed to set self.output.{batch_dim_axis,time_dim_axis} explicitly,
    # as well as check the inputs if they are as you would suggest.
    # However, a good default is often to use the same as the input.
    if all([k not in out_type for k in Data.SpecialAxesNames]):
      if sources_data:
        out_type.setdefault("batch_dim_axis", sources_data.batch_dim_axis)
        out_type.setdefault("time_dim_axis", sources_data.time_dim_axis)
        if not out_type.get("sparse", False) and sources_data.feature_dim_axis_or_unspecified is not NotSpecified:
          if sources_data.feature_dim_axis_or_unspecified is not None:
            out_type.setdefault("feature_dim_axis", sources_data.feature_dim_axis_or_unspecified)
          else:  # None
            if out_type.get("dim", None) is None:
              out_type.setdefault("feature_dim_axis", None)
      elif network.is_inside_rec_layer() and None not in out_type.get("shape", ()):
        out_type.setdefault("time_dim_axis", None)
    if "shape" not in out_type:
      if sources_data:
        if out_type.get("sparse", False):
          out_type.setdefault("shape", sources_data.shape_sparse)
        else:  # not sparse
          feature_dim_axis = out_type.get("feature_dim_axis", NotSpecified)
          if feature_dim_axis is NotSpecified:
            if sources_data.feature_dim_axis is not None:
              feature_dim_axis = sources_data.feature_dim_axis
            else:
              feature_dim_axis = -1
          if sources_data.shape:
            default_shape = list(sources_data.shape_dense)
            if sources_data.batch_dim_axis is not None:
              default_shape.insert(sources_data.batch_dim_axis, None)
            default_shape[feature_dim_axis] = out_type.get("dim", None)
            if out_type.get("batch_dim_axis") is not None:
              default_shape.pop(out_type.get("batch_dim_axis"))
          else:  # source is scalar
            if out_type.get("dim") or out_type.get("feature_dim_axis") is not None:
              default_shape = (out_type.get("dim"),)
            else:
              default_shape = ()
          out_type.setdefault("shape", tuple(default_shape))
      elif network.is_inside_rec_layer():
        if out_type.get("sparse", False):
          out_type.setdefault("shape", ())
        else:
          out_type.setdefault("shape", (out_type.get("dim", None),))
    # Note: No special handling for feature_dim_axis here for now...
    beam = None
    for src in sources:
      if src:  # might be None if template construction
        beam = SearchBeam.get_combined_beam(beam, src.output.beam)
    out_type.setdefault("beam", beam)
    output = Data(**out_type)
    cls._post_init_output(
      output=output, network=network, target=target, size_target=size_target, _target_layers=_target_layers,
      sources=sources, **kwargs)
    return output

  # noinspection PyUnusedLocal
  @classmethod
  def _post_init_output(cls, output, network, target=None, size_target=None, _target_layers=None, sources=(),
                        _src_common_search_choices=None,
                        **kwargs):
    """
    :param Data output:
    :param returnn.tf.network.TFNetwork network:
    :param str|list[str]|None target:
    :param str|None size_target:
    :param dict[str,LayerBase]|None _target_layers: if target.startswith("layer:"), then this is target -> layer
    :param None|SearchChoices _src_common_search_choices: set via :func:`SearchChoices.translate_to_common_search_beam`
    :param list[LayerBase] sources:
    """
    # You are supposed to set self.output.placeholder to the value which you want to return by the layer.
    # Normally you are also supposed to set self.output.size_placeholder explicitly, just like self.output.placeholder.
    # However, in many cases, this will just be {0: time-lengths} and the same as from the input.
    # We check for this case and preset it by that if possible.
    # If you want to have it different in your layer, just overwrite it.
    common_source = Data.get_common_data([s.output for s in sources if s])
    if not output.size_placeholder:
      if network.eval_flag and size_target:
        output.size_placeholder = cls._static_get_target_value(
          target=size_target,
          _target_layers=_target_layers,
          search_choices=_src_common_search_choices,
          network=network, mark_data_key_as_used=network.eval_flag).size_placeholder.copy()
      elif common_source and common_source.matches_var_dim_pattern(output):
        output.size_placeholder = common_source.size_placeholder.copy()
      elif network.train_flag is not False and target:
        # TODO: In training, this is ok. Maybe as well as for eval but not clear.
        # In forward, mark_data_key_as_used=False should be used and anyway that target value is not available.
        output.size_placeholder = cls._static_get_target_value(
          target=(target[0] if (target and isinstance(target, list)) else target),
          _target_layers=_target_layers,
          search_choices=_src_common_search_choices,
          network=network, mark_data_key_as_used=network.train_flag is not False).size_placeholder.copy()
    if any([(src and not src.output.available_for_inference) for src in sources if src]):
      output.available_for_inference = False

  @classmethod
  def cls_get_tf_scope_name(cls, name):
    """
    :param str name: layer name
    :return: valid scope name, might be just name. see tf._VALID_SCOPE_NAME_REGEX and tf._VALID_OP_NAME_REGEX
    :rtype: str
    """
    from returnn.tf.util.basic import get_valid_scope_name_from_str
    return get_valid_scope_name_from_str(name)

  @classmethod
  def cls_layer_scope(cls, name):
    """
    Setup scope for layer. This can also be used when the layer does not yet exists.
    This is supposed to cover variable creations as well.
    Currently vars might be created when used within the rec-layer, but they are caught
    in a more generic way there, so we have not implemented yet any special logic here.

    :param str name: layer name
    :return: context manager object
    """
    @contextlib.contextmanager
    def layer_scope_ctx():
      """
      :return: context manager object
      """
      from returnn.tf.util.basic import reuse_name_scope
      with reuse_name_scope(cls.cls_get_tf_scope_name(name)) as scope:
        yield scope
    return layer_scope_ctx()

  @classmethod
  def get_global_layer_list(cls):
    """
    :rtype: list[LayerBase]
    """
    from returnn.tf.util.basic import CollectionKeys
    coll = tf_compat.v1.get_collection_ref(CollectionKeys.RETURNN_LAYERS)
    assert isinstance(coll, list)
    return coll

  @classmethod
  def get_recent_layer(cls):
    """
    :rtype: LayerBase
    """
    coll = cls.get_global_layer_list()
    assert coll
    return coll[-1]

  def _register_layer(self):
    self.get_global_layer_list().append(self)

  @classmethod
  def transform_config_dict(cls, d, network, get_layer):
    """
    :param dict[str] d: will modify inplace
    :param returnn.tf.network.TFNetwork network:
    :param ((str) -> LayerBase) get_layer: function to get or construct another layer
      The name `get_layer` might be misleading, as this should return an existing layer,
      or construct it if it does not exist yet.
      `network.get_layer` would just return an existing layer.

    Will modify `d` inplace such that it becomes the kwargs for `self.__init__()`.
    Mostly leaves `d` as-is.
    This is used by :func:`TFNetwork.construct_from_dict`.
    It resolves certain arguments,
    e.g. it resolves the `"from"` argument which is a list of strings,
    to make it the `"sources"` argument in kwargs, with a list of :class:`LayerBase` instances.
    Subclasses can extend/overwrite this.
    Usually the only reason to overwrite this is when some argument might be a reference to a layer
    which should be resolved.
    """
    from .basic import get_loss_class
    src_names = d.pop("from", ["data"])
    if not isinstance(src_names, (list, tuple)):
      src_names = [src_names]
    d["sources"] = [
      get_layer(src_name)
      for src_name in src_names
      if not src_name == "none"]
    if "collocate_with" in d:
      collocate_with = d["collocate_with"]
      if not isinstance(collocate_with, (list, tuple)):
        collocate_with = [collocate_with]
      d["collocate_with"] = collocate_with  # not get_layer: we don't really want this dependency
    if "reuse_params" in d:
      d["reuse_params"] = ReuseParams.from_config_dict(d["reuse_params"], network=network, get_layer=get_layer)
    if d.get("loss", None) and "target" not in d:
      target = get_loss_class(d["loss"]).get_default_target(network.extern_data)
      if target:
        d["target"] = target
    targets = None
    target_layers = {}
    assert "_target_layers" not in d
    if d.get("target", None):
      targets = d["target"]
      # we might have multiple targets, e.g. in choice layer, so convert to list
      if isinstance(targets, str):
        targets = [targets]
      if network.eval_flag:
        # _target_layers is a small workaround for further code which might not have access to the right get_layer.
        d["_target_layers"] = target_layers
        for target in targets:
          assert isinstance(target, str)
          # Not resolving this in the dict as target, but call get_layer to make it available.
          if target.startswith("layer:"):
            target_layers[target] = get_layer(target[len("layer:"):])
          else:
            # Note: This is a workaround for cases where we need to know about used data keys before the layer
            # itself is constructed (e.g. in _SubnetworkRecCell.get_output).
            # A nicer solution would be to not modify this here,
            # but instead lazily handle it in TFNetwork.get_extern_data,
            # such that we do not need to know in advance which data keys we need.
            # Also, if we are inside a rec layer, and doing search, we also cannot do that.
            if network.is_inside_rec_layer() and not network.search_flag:
              network.get_extern_data(target, mark_data_key_as_used=True)
            if not network.search_flag:
              # Also, there are cases when we want to have the target as an explicit layer dep,
              # e.g. when the target has a beam, to derive the search choices.
              target_layers[target] = get_layer("data:%s" % target)
    if d.get("initial_output", None):  # see get_rec_initial_output
      initial_output = d["initial_output"]
      if isinstance(initial_output, str):
        if initial_output not in ["zeros", "ones", "var", "keep_over_epoch", "keep_over_epoch_no_init", "apply(0)"]:
          # if initial_output is not a reserved keyword, assume it is a layer
          d['initial_output'] = get_layer(initial_output)
    if "n_out" not in d and targets:
      # Must be done here now because loss might be set to None later.
      target = targets[0]  # guess using first target
      d["n_out"] = cls._guess_n_out_from_target_and_opt_loss(
        network=network, target=target, target_layers=target_layers,
        loss_class_name=d.get("loss", None), get_layer=get_layer)
    if d.pop("loss_only_on_non_search", None) and network.search_flag:
      d.pop("loss", None)
      d.pop("loss_scale", None)
      d.pop("loss_opts", None)
    if d.get("loss", None):
      loss_opts = d.pop("loss_opts", None)
      if not loss_opts:
        loss_opts = {}
      # loss_scale: scale factor for loss (1.0 by default). DEPRECATED: use loss.scale instead, via loss_opts
      loss_scale = d.pop("loss_scale", 1.0)
      if loss_scale != 1.0:
        assert loss_opts.get("scale", 1.0) == 1.0, "do not use loss_scale and loss with 'scale' option together"
        loss_opts["scale"] = loss_scale
      d["loss"] = cls._make_loss(
        class_name=d.pop("loss", None), opts=loss_opts, network=network, get_layer=get_layer)
    else:
      if "loss_scale" in d and d["loss_scale"] is None:
        d.pop("loss_scale")
      if "loss_opts" in d and d["loss_opts"] is None:
        d.pop("loss_opts")
      assert "loss_scale" not in d, "loss not defined, do not set loss_scale"
      assert "loss_opts" not in d, "loss not defined, do not set loss_opts"

  @classmethod
  def _guess_n_out_from_target_and_opt_loss(cls, network, target, target_layers, loss_class_name, get_layer):
    """
    :param returnn.tf.network.TFNetwork network:
    :param str target: e.g. "classes"
    :param dict[str,LayerBase] target_layers:
    :param str|None loss_class_name: e.g. "ce" or None
    :param ((str) -> LayerBase) get_layer: function to get or construct another layer
    :return: n_out value
    :rtype: int
    """
    from .basic import get_loss_class
    if target in target_layers:
      target_data = target_layers[target].output
    else:
      target_data = cls._static_get_target_value(
        target=target, network=network, mark_data_key_as_used=False, get_layer=get_layer)
      if not target_data:
        return 1  # dummy value during template construction. this would be corrected later
    n_out = target_data.dim
    if loss_class_name:
      n_out = get_loss_class(loss_class_name).get_auto_output_layer_dim(n_out)
    return n_out

  @classmethod
  def _make_loss(cls, class_name, opts, network, get_layer, always_make=False):
    """
    :param str|None class_name:
    :param dict[str]|None opts:
    :param returnn.tf.network.TFNetwork network:
    :param ((str) -> LayerBase) get_layer: function to get or construct another layer
    :param bool always_make:
    :rtype: Loss|None
    """
    from .basic import get_loss_class
    if not network.eval_flag and not always_make:
      # Don't resolve the loss opts on purpose.
      # This might result in a smaller network because it might skip some get_layer calls.
      # This is what we want, i.e. we don't want to resolve layers which are only needed for the loss.
      return None
    if not class_name:
      assert not always_make
      return None
    if not opts:
      opts = {}
    opts = opts.copy()
    loss_class = get_loss_class(class_name)
    assert issubclass(loss_class, Loss)
    loss_class.transform_config_dict(opts, network=network, get_layer=get_layer)
    loss = loss_class(base_network=network, **opts)
    assert isinstance(loss, Loss)
    return loss

  @property
  def tf_scope_name(self):
    """
    :rtype: str
    :return: normally just self.name, but make it a valid TF scope name
    """
    return self.cls_get_tf_scope_name(name=self.name)

  def get_base_absolute_name_scope_prefix(self):
    """
    :return: e.g. "output/", always with "/" at end
    :rtype: str
    """
    return self.network.get_absolute_name_scope_prefix() + self.tf_scope_name + "/"

  def get_absolute_name_scope_prefix(self):
    """
    :return: e.g. "output/", always with "/" at end
    :rtype: str
    """
    return self.get_base_absolute_name_scope_prefix()

  def get_absolute_name(self):
    """
    :return: e.g. "output" or "subnet/output"
    :rtype: str
    """
    return self.network.get_absolute_name_prefix() + self.name

  def is_output_layer(self):
    """
    Some code differs between an output layer and other layers.
    It is a bit arbitrary what we define as output layer.
    This should be consistent with :func:`TFNetwork.construct_from_dict`.

    :rtype: bool
    """
    if self._is_output_layer is not None:
      return self._is_output_layer
    if self.loss:
      return True
    if self.name == "output":
      return True
    return False

  def get_dep_layers(self):
    """
    :return: list of layers this layer depends on.
      normally this is just self.sources but e.g. the attention layer in addition has a base, etc.
    :rtype: list[LayerBase]
    """
    layers = list(self.sources)
    if self._target_layers:
      layers += [layer for _, layer in sorted(self._target_layers.items())]
    return layers

  def get_sub_layer(self, layer_name):
    """
    The default behavior for any layer is to return None.

    :param str layer_name: name of the sub_layer (right part of '/' separated path)
    :return: the sub_layer addressed in layer_name or None if no sub_layer exists
    :rtype: LayerBase|None
    """
    return None

  @classmethod
  def get_sub_layer_out_data_from_opts(cls, layer_name, parent_layer_kwargs):
    """
    Called by _TemplateLayer.get_sub_layer(). Gets a Data template for the sub-layer with name 'layer_name'.
    Also returns the network the sub-layer is in and the class type of the sub-layer. There is no good default
    behaviour here, as this heavily depends on how the current layer uses sub-layers.

    :param str layer_name: name of the sub_layer (right part of '/' separated path)
    :param dict[str] parent_layer_kwargs: kwargs for the parent layer (as kwargs in cls.get_out_data_from_opts())
    :return: Data template, network and the class type of the sub-layer
    :rtype: (Data, TFNetwork, type)|None
    """
    return None

  def get_search_choices(self):
    """
    :rtype: SearchChoices|None
    """
    if self.search_choices:
      return self.search_choices
    if self._src_common_search_choices:
      return self._src_common_search_choices
    if not self.output.beam:  # small optimization
      # Note that the logic to determine self.output.beam (via get_out_data_from_opts) is currently
      # independently implemented from the search choices (via get_search_choices),
      # and due to bugs, it could happen that self.output.beam is not set but actually we have search choices.
      # This will likely be caught at some later check.
      return None
    layer = self.network.get_search_choices(src=self)
    if layer:
      assert layer.search_choices
      return layer.search_choices
    return None

  def get_search_beam_size(self):
    """
    :return: beam size if there was a choice layer and we do search
    :rtype: int|None
    """
    if self.output.beam:
      return self.output.beam.beam_size
    return None

  def get_normalized_layer(self):
    """
    :return: e.g. if prev layer in :class:`RecLayer`, return current layer
    :rtype: LayerBase
    """
    return self

  def get_batch_dim(self):
    """
    The batch dim by this layer, not taken from our output but calculated.
    Normally it is self.network.get_batch_dim()
    but if we do search and there was a choice layer, it it multiplied by the beam size.
    :return: batch dim * beam size
    :rtype: tf.Tensor|int
    """
    batch_dim = self.network.get_data_batch_dim()
    beam_size = self.get_search_beam_size()
    if beam_size is not None:
      with tf.name_scope("batch_beam_dim"):
        batch_dim *= beam_size
    return batch_dim

  @contextlib.contextmanager
  def var_creation_scope(self, **kwargs):
    """
    This takes care of setting up a scope where variables can be created.
    This handles multiple things:

     * the param sharing logic, to reuse existing variables from elsewhere
     * variational noise
     * Note: :func:`default_control_flow_ctx` should not be needed, as tf.get_variable should always work

    :param kwargs: passed to variable_scope
    :return: yields the variable_scope
    """
    from returnn.tf.util.basic import get_current_var_scope_name, reuse_name_scope
    from returnn.tf.util.basic import default_control_flow_ctx, reuse_name_scope_of_tensor
    self_base_scope = self.get_base_absolute_name_scope_prefix()
    assert self_base_scope.endswith("/")
    cur_scope = get_current_var_scope_name()
    assert (cur_scope + "/").startswith(self_base_scope)
    # There are cases were a dummy layer was created already to create the variables,
    # e.g. see ReuseParams.LazyLayerResolver.
    kwargs = kwargs.copy()
    kwargs.setdefault("reuse", getattr(tf_compat.v1, "AUTO_REUSE", None))

    param_variational_noise = self.param_variational_noise
    if param_variational_noise is None:
      param_variational_noise = self.network.get_config().float("param_variational_noise", 0)
    if self.network.train_flag is False:  # if True or tf.Tensor, it will use cond_on_train below
      param_variational_noise = None
    need_custom_getter = bool(param_variational_noise)  # and param.dtype.is_floating
    kwargs_custom_getter = kwargs.get("custom_getter", None)

    def layer_custom_getter(getter, **getter_kwargs):
      """
      See TF docs :func:`_VariableStore.get_variable`.

      :param (...)->tf.Variable getter:
      :rtype: tf.Variable|tf.Tensor
      """
      if kwargs_custom_getter:
        param = kwargs_custom_getter(getter, **getter_kwargs)
      else:
        param = getter(**getter_kwargs)

      # Only apply this if we get a variable. Otherwise, maybe variational noise was already applied
      # (by some parent var scope), and we don't want to apply it twice.
      if param_variational_noise and param.dtype.is_floating and isinstance(param, tf.Variable):
        with default_control_flow_ctx():  # make independent from loop/cond
          with reuse_name_scope_of_tensor(param, postfix="_variational_noise", add_tensor_name=True):
            param = self.network.cond_on_train(
              fn_train=lambda: param + tf_compat.v1.random_normal(
                tf.shape(param), dtype=param.dtype.base_dtype,
                stddev=param_variational_noise,
                seed=self.network.random.randint(2 ** 31)),
              fn_eval=lambda: param)

      return param

    @contextlib.contextmanager
    def inner():
      """
      Var creation scope + variable scope.
      """
      if self.reuse_params:
        var_scope = self.reuse_params.get_variable_scope(base_layer=self)
      else:
        var_scope = tf_compat.v1.get_variable_scope()
      if need_custom_getter:
        kwargs["custom_getter"] = layer_custom_getter
      with reuse_name_scope(var_scope, **kwargs) as scope_:
        yield scope_

    if self.param_device:
      device_name = self.param_device
      if ":" not in device_name:
        device_name = "%s:*" % device_name
      if "/" not in device_name:
        device_name = "/device:%s" % device_name
      with tf.device(device_name):
        with inner() as scope:
          yield scope
    else:
      with inner() as scope:
        yield scope

  def add_param(self, param, custom_update=None, trainable=None, saveable=None, axes_split_info=None):
    """
    :param tf.Variable|tf.Tensor param:
    :param None|CustomUpdate custom_update: will be applied in training, instead of taking the gradient
    :param bool|None trainable:
    :param bool|None saveable:
    :param list[list[int]]|None axes_split_info: e.g. [[n],[n]*4] for LSTM matrices
    :return: param
    :rtype tf.Variable
    """
    _param = param
    if isinstance(param, tf.Tensor):
      # This can happen with a custom_getter in tf.compat.v1.get_variable(), e.g. via self.reuse_params.
      # Check if we can still find the original variable.
      from returnn.extern import graph_editor
      import re
      possible_params = tf_compat.v1.get_collection(
        tf_compat.v1.GraphKeys.GLOBAL_VARIABLES, scope=re.escape(self.get_absolute_name_scope_prefix()))
      if not possible_params:
        # None found. Just return as-is.
        return param
      all_ops = graph_editor.get_backward_walk_ops([param.op], inclusive=False, control_inputs=False)
      all_1st_tensors = [op.outputs[0] for op in all_ops if len(op.outputs) == 1]
      # noinspection PyProtectedMember
      possible_params = [p for p in possible_params if tf_util.var_handle_or_ref(p) in all_1st_tensors]
      if not possible_params:
        # Not found. Just return as-is.
        return param
      assert len(possible_params) == 1
      param = possible_params[0]
    assert isinstance(param, tf.Variable)
    if not self.trainable:
      trainable_collection_ref = param.graph.get_collection_ref(tf_compat.v1.GraphKeys.TRAINABLE_VARIABLES)
      if param in trainable_collection_ref:
        trainable_collection_ref.remove(param)
    if trainable is None:
      trainable = param in param.graph.get_collection_ref(tf_compat.v1.GraphKeys.TRAINABLE_VARIABLES)
    if saveable is None:
      saveable = True
    if custom_update:
      assert trainable
      custom_update.set_on_var(param)
    if axes_split_info:
      tf_util.set_param_axes_split_info(param, axes_split_info)
    if self.reuse_params:
      name_scope_prefix = self.reuse_params.get_absolute_name_scope_prefix(base_layer=self, param=param)
    else:
      name_scope_prefix = self.get_absolute_name_scope_prefix()
    assert param.name
    assert param.name[:len(name_scope_prefix)] == name_scope_prefix
    assert param.name[-2:] == ":0"
    param_name = param.name[len(name_scope_prefix):-2]
    if param_name not in self.params:
      self.params[param_name] = param
    else:
      assert self.params[param_name] is param
    if not saveable:
      self.saveable_param_replace[param] = None
    if getattr(param, "RETURNN_layer", None) is None:
      param.RETURNN_layer = self
    if getattr(param, "RETURNN_updater_opts", None) is None and self.updater_opts.truth_value:
      param.RETURNN_updater_opts = self.updater_opts
    # Note that any further postprocessing on the parameter should not be done here,
    # as we cannot guarantee that the result from this method is really used,
    # e.g. when we use official TF code such as the official LSTM cell.
    # The better way is to do it in self.var_creation_scope(), which also applies in those cases.
    return _param

  def set_param_values_by_dict(self, values_dict, session, ignore_wrong_shape=False, copy_param_mode=None):
    """
    :param dict[str,numpy.ndarray] values_dict:
    :param bool ignore_wrong_shape:
    :param str|None copy_param_mode:
    :param tf.compat.v1.Session session:
    """
    if callable(self.custom_param_importer):
      self.custom_param_importer(layer=self, values_dict=values_dict, session=session)
      return
    if self.custom_param_importer:
      copy_param_mode = self.custom_param_importer
    if copy_param_mode == "reset":
      return  # just ignore. the random init was already done
    assert copy_param_mode in [None, "ifpossible", "subset"]
    if copy_param_mode:
      ignore_wrong_shape = True
    for param_name, values in values_dict.items():
      assert param_name in self.params, '%s: param %r unknown' % (self, param_name)
      param = self.params[param_name]
      assert isinstance(param, tf.Variable)
      shape = param.get_shape()
      assert isinstance(shape, tf.TensorShape)
      assert shape.is_fully_defined(), '%s: shape of param %r %r not fully defined?' % (self, param_name, param)
      param_shape = tuple(shape.as_list())
      if not ignore_wrong_shape:
        assert param_shape == values.shape, "var %r: shape %s != %s" % (param, shape.as_list(), values.shape)
      if param_shape != values.shape:
        if copy_param_mode == "subset":
          assert len(param_shape) == len(values.shape), "param %r ndim must match" % param
          new_values = session.run(param)  # use currently (randomly) initialized params as base
          param_axes_split_info = tf_util.get_param_axes_split_info(param)
          if param_axes_split_info:
            tf_util.check_param_axes_split_info(param.get_shape().as_list(), param_axes_split_info)
            old_axes_splits = tf_util.transform_param_axes_split_info_to_new_shape(
              param_axes_split_info, values.shape, debug_name="param %r" % param.name)
            print("Param %r: transform old values of shape parts %r into new shape parts %r." % (
              param, old_axes_splits, param_axes_split_info), file=log.v3)
            values = tf_util.copy_with_new_split_axes(
              old_axis_splits=old_axes_splits, new_axis_splits=param_axes_split_info,
              old_values=values, new_values=new_values)
          else:
            print("Param %r: transform old values of shape %r into new shape %r." % (
              param, values.shape, param_shape), file=log.v3)
            values = tf_util.copy_with_new_split_axes(
              old_axis_splits=[[d] for d in values.shape],
              new_axis_splits=[[d] for d in param_shape],
              old_values=values, new_values=new_values)
        else:
          print(
            "Will not set param %r because its shape %s != %s." % (param, shape.as_list(), values.shape), file=log.v3)
          continue
      self.network.get_var_assigner(param).assign(values, session=session)

  def get_param_values_dict(self, session):
    """
    :param tf.compat.v1.Session session:
    :return: dict name -> values
    :rtype: dict[str,numpy.ndarray]
    """
    d = {}
    for param_name, param in self.get_saveable_params_dict().items():
      d[param_name] = param.eval(session)
    return d

  def get_saveable_params_dict(self):
    """
    :return: params and saveable_param_replace resolved
    :rtype: dict[str,tf.Variable|tensorflow.python.training.saver.BaseSaverBuilder.SaveableObject]
    """
    if not self.saveable_param_replace:
      return self.params.copy()
    d = {}
    for param_name, param in self.params.items():
      if param in self.saveable_param_replace:
        param = self.saveable_param_replace[param]
        if param is None:
          continue
      d[param_name] = param
    return d

  @staticmethod
  def _static_get_target_value(
    target, network, mark_data_key_as_used=True,
    _target_layers=None, get_layer=None,
    search_choices=None
  ):
    """
    :param str target:
    :param dict[str,LayerBase]|None _target_layers: if target.startswith("layer:"), then this is target -> layer
    :param returnn.tf.network.TFNetwork network:
    :param bool mark_data_key_as_used: forwarded self.network.get_extern_data()
    :param None|((str) -> LayerBase) get_layer: function to get or construct another layer
    :param SearchChoices|None search_choices:
    :rtype: Data | None
    """
    if not target or target == "none":
      return None
    from .basic import SelectSearchSourcesLayer
    if _target_layers and target in _target_layers:
      return SelectSearchSourcesLayer.select_if_needed(
        _target_layers[target], search_choices=search_choices).output
    if target.startswith("layer:"):
      if not get_layer:
        get_layer = network.get_layer
      layer = get_layer(target[len("layer:"):])
      if not layer:  # some get_layer during temp construction might return None
        return None
      return SelectSearchSourcesLayer.select_if_needed(layer, search_choices=search_choices).output
    assert network.extern_data.has_data(target), "target %r unknown" % target
    data = network.get_extern_data(target, mark_data_key_as_used=mark_data_key_as_used)
    if search_choices:
      data = data.copy_extend_with_beam(search_choices.get_beam_info())
    return data

  def _get_target_value(self, target=None, mark_data_key_as_used=True, search_choices=NotSpecified):
    """
    :param str|None target:
    :param bool mark_data_key_as_used: forwarded self.network.get_extern_data()
    :param SearchChoices|NotSpecified|None search_choices:
    :rtype: Data | None
    """
    if target is None:
      target = self.target
    if search_choices is NotSpecified:
      search_choices = self.get_search_choices()
    return self._static_get_target_value(
      target=target, _target_layers=self._target_layers, search_choices=search_choices,
      network=self.network, mark_data_key_as_used=mark_data_key_as_used)

  def _cond_only_on_eval_opt(self, on_eval_func, default_value):
    """
    :param ()->(tf.Tensor|None) on_eval_func:
    :param float|tf.Tensor default_value:
    :return: tensor (coming from tf.cond if needed) if on_eval_func returned a tensor, otherwise None
    :rtype: tf.Tensor|None
    """
    if not isinstance(default_value, tf.Tensor):
      default_value = tf.constant(default_value, name="only_on_eval_dummy_zero")

    class OnEval:
      """
      Closure.
      """
      have_output = True

      @classmethod
      def get_value(cls):
        """
        :rtype: tf.Tensor
        """
        res_ = on_eval_func()
        if res_ is None:
          cls.have_output = False
          return default_value  # Doesn't matter, will not be used anyway.
        return res_

    res = self.network.cond_on_train(lambda: default_value, OnEval.get_value)
    if not OnEval.have_output:
      return None
    return res

  @classmethod
  def get_losses(cls, name, network, output, loss=None, reduce_func=None, layer=None, **kwargs):
    """
    Losses will get constructed here.
    This gets called inside a loss name scope of the layer.
    When overriding this, make sure that it works both with `layer` set and unset.

    :param str name: layer name
    :param returnn.tf.network.TFNetwork network:
    :param Loss|None loss: argument just as for __init__
    :param Data output: the output (template) for the layer
    :param LayerBase|None layer:
      The real layer instance, if it exists at the current point.
      If not given, init() must be called at a later point.
    :param ((tf.Tensor)->tf.Tensor)|None reduce_func: if given, will overwrite the reduce func for the loss.
      By default, every loss_value and error_value is a scalar
      (sum or average over the batches, and over the frames for frame-wise losses).
      However, if you provide reduce_func = TFUtil.identity, you can get the unreduced tensor.
    :param kwargs: all the remaining __init__ args
    :return: the losses defined by this layer
    :rtype: list[returnn.tf.network.LossHolder]
    """
    if not loss:
      return []
    from returnn.tf.network import LossHolder
    return [LossHolder(
      name=name, network=network, loss=loss, layer_output=output, layer=layer, reduce_func=reduce_func)]

  def get_losses_initialized(self, reduce_func=None):
    """
    As self.get_losses, but here we return them all initialized (i.e. the layer is set).
    You should not override this method but rather :func:`get_losses`.

    :param ((tf.Tensor)->tf.Tensor)|None reduce_func: as in get_losses
    :return: the losses defined by this layer
    :rtype: list[returnn.tf.network.LossHolder]
    """
    return self.__class__.get_losses(reduce_func=reduce_func, layer=self, **self.kwargs)

  def get_params_l2_norm(self):
    """
    :return: scalar
    :rtype: tf.Tensor
    """
    return 2 * sum([tf.nn.l2_loss(param) for (name, param) in sorted(self.params.items())])

  def get_output_spatial_smoothing_energy(self):
    """
    :return: scalar. see :func:`TFUtil.spatial_smoothing_energy`
    :rtype: tf.Tensor
    """
    from returnn.tf.util.basic import spatial_smoothing_energy, flatten_with_seq_len_mask
    energy = spatial_smoothing_energy(self.output.placeholder, dim=self.output.dim)  # (batch,time)
    assert self.output.have_time_axis()
    energy = flatten_with_seq_len_mask(
      energy,
      seq_lens=self.output.size_placeholder[self.output.time_dim_axis_excluding_batch],
      time_major=self.output.is_time_major)  # (time')
    energy = tf.reduce_sum(energy)
    return energy

  def get_darc1(self):
    """
    DARC1, simplified Directly Approximately Regularizing Complexity (DARC), via
    Generalization in Deep Learning, https://arxiv.org/abs/1710.05468

    :return: scalar
    :rtype: tf.Tensor
    """
    with tf.name_scope("darc1"):
      if self.output_before_activation:
        x = self.output_before_activation.x
      else:
        x = self.output.placeholder
      mask = self.output.get_sequence_mask()  # (time,batch) or (batch,time), like output
      size = tf.size(mask)  # time * batch
      mask = tf.reshape(mask, (size,))  # (time*batch,)
      x = tf.reshape(x, (size,) + self.output.shape[1:])  # (time*batch,dim)
      x = tf.abs(x)
      x = tf.where(mask, x, tf.zeros_like(x))
      x = tf.reduce_sum(x, axis=0)  # (dim,)
      assert isinstance(x, tf.Tensor)
      assert x.get_shape().ndims == 1
      x = tf.reduce_max(x)  # scalar
      return x

  def get_constraints_value(self):
    """
    :return: None or scalar
    :rtype: tf.Tensor|None
    """
    c = 0
    if self.L2:
      c += self.L2 * self.get_params_l2_norm()
    if self.spatial_smoothing:
      c += self.spatial_smoothing * self.get_output_spatial_smoothing_energy()
    if self.darc1:
      c += self.darc1 * self.get_darc1()
    if c is 0:
      return None
    return c

  def batch_norm(self, data,
                 use_shift=True, use_std=True, use_sample=0.0, force_sample=False,
                 momentum=0.99, epsilon=1e-3,
                 sample_mean=None, sample_variance=None,
                 update_sample_only_in_training=False,
                 delay_sample_update=False,
                 param_version=0,
                 gamma=None, beta=None,
                 gamma_init=1.0, beta_init=0.0,
                 masked_time=True):
    """
    :param Data data:
    :param bool use_shift:
    :param bool use_std:
    :param float use_sample: defaults to 0.0 which is used in training
    :param bool force_sample: even in eval, use the use_sample factor
    :param float momentum: for the running average of sample_mean and sample_std
    :param bool update_sample_only_in_training:
    :param bool delay_sample_update:
    :param int param_version: 0 or 1
    :param float epsilon:
    :param tf.Tensor sample_mean:
    :param tf.Tensor sample_variance:
    :param tf.Tensor gamma:
    :param tf.Tensor beta:
    :param str|float gamma_init: see :func:`TFUtil.get_initializer`, for the scale
    :param str|float beta_init: see :func:`TFUtil.get_initializer`, for the mean
    :param bool masked_time: flatten and mask input tensor
    :rtype: tf.Tensor

    http://arxiv.org/abs/1502.03167

    With our default settings:

    - In training: use_sample=0, i.e. not using running average, using current batch mean/var.
    - Not in training (e.g. eval): use_sample=1, i.e. using running average, not using current batch mean/var.
    - The running average includes the statistics of the current batch.
    - The running average is also updated when not training.

    Also see:
      tf.nn.batch_normalization()
      https://github.com/deepmind/sonnet/blob/master/sonnet/python/modules/batch_norm.py
    """
    with reuse_name_scope(self.get_absolute_name_scope_prefix() + "batch_norm", absolute=True):
      if masked_time:
        x = data.get_placeholder_flattened(keepdims=True)
        mean, variance = tf_compat.v1.nn.moments(x, axes=[0], keep_dims=True)
      else:
        x = data.placeholder
        mean, variance = tf_compat.v1.nn.moments(x, axes=data.get_axes(exclude_feature=True), keep_dims=True)
      if update_sample_only_in_training:
        momentum = tf.where(self.network.train_flag, momentum, 0.0)
      delayed_ops = []  # type: typing.List[tf.Operation]
      # Note about param_version:
      # We might later drop some earlier versions.
      # We just need to implement a conversion in CustomCheckpointLoader.
      if param_version == 0:
        param_name_prefix = "%s_%s_" % (self.name, data.name)
      elif param_version == 1:
        param_name_prefix = ""
      else:
        raise NotImplementedError("%s: batch_norm param_version %r" % (self, param_version))
      if sample_mean is None:
        with self.var_creation_scope():
          sample_mean = self.add_param(tf_compat.v1.get_variable(
            shape=data.get_bc_spatial_batch_shape(), initializer=tf_compat.v1.zeros_initializer(),
            name="%smean" % param_name_prefix,
            trainable=False))
        # Use exponential moving average of batch mean.
        # Note: We could also use cumulative moving average. Our Theano implementation does that for inference.
        updated_sample_mean = tf_compat.v1.assign_add(sample_mean, (mean - sample_mean) * momentum)
        if delay_sample_update:
          delayed_ops.append(updated_sample_mean.op)
        else:
          sample_mean = updated_sample_mean
      if sample_variance is None:
        # Note: Our Theano implementation does not use a moving average for this.
        with self.var_creation_scope():
          sample_variance = self.add_param(tf_compat.v1.get_variable(
            shape=data.get_bc_spatial_batch_shape(), initializer=tf_compat.v1.ones_initializer(),
            name="%svariance" % param_name_prefix,
            trainable=False))
        updated_sample_variance = tf_compat.v1.assign_add(sample_variance, (variance - sample_variance) * momentum)
        if delay_sample_update:
          delayed_ops.append(updated_sample_variance.op)
        else:
          sample_variance = updated_sample_variance
      # If train or if force_sample, use default use_sample=0.0, otherwise use_sample=1.0.
      if self.network.train_flag is not False or force_sample:
        if force_sample:
          pass  # leave use_sample as-is
        else:
          use_sample = tf.where(self.network.train_flag, use_sample, 1.0)
      else:
        use_sample = 1.0
      mean = (1. - use_sample) * mean + use_sample * sample_mean
      variance = (1. - use_sample) * variance + use_sample * sample_variance
      bn = (data.placeholder - mean) * tf_compat.v1.rsqrt(variance + epsilon)
      if delayed_ops:
        for op in delayed_ops:
          # Make sure we update after we calculated the batch norm.
          tf_util.add_control_input(op, control_input=bn.op)
        self.network.register_post_control_dependencies(delayed_ops)
      if use_std:
        if gamma is None:
          with self.var_creation_scope():
            from returnn.tf.util.basic import get_initializer
            gamma_initializer = get_initializer(
              gamma_init, seed=self.network.random.randint(2 ** 31) if gamma_init else 0, eval_local_ns={"layer": self})
            gamma = self.add_param(tf_compat.v1.get_variable(
              shape=data.get_bc_spatial_batch_shape(), initializer=gamma_initializer,
              name="%sgamma" % param_name_prefix,
              trainable=True))
        bn *= gamma
      if use_shift:
        if beta is None:
          with self.var_creation_scope():
            from returnn.tf.util.basic import get_initializer
            beta_initializer = get_initializer(
              beta_init, seed=self.network.random.randint(2 ** 31) if beta_init else 0, eval_local_ns={"layer": self})
            beta = self.add_param(tf_compat.v1.get_variable(
              shape=data.get_bc_spatial_batch_shape(), initializer=beta_initializer,
              name="%sbeta" % param_name_prefix,
              trainable=True))
        bn += beta
      return bn

  def get_hidden_state(self):
    """
    If this is a recurrent layer, this would return the hidden state.
    This is used e.g. for the RnnCellLayer class.

    :rtype: tf.Tensor | list[tf.Tensor] | None
    :return: optional tensor(s) with shape (time, batch, dim)
    """
    return None

  def get_last_hidden_state(self, key):
    """
    If this is a recurrent layer, this would return the last hidden state.
    Otherwise, we return None.

    :param int|str|None key: also the special key "*"
    :rtype: tf.Tensor | None
    :return: optional tensor with shape (batch, dim)
    """
    if key in self.rec_vars_outputs:
      return self.rec_vars_outputs[key]
    if key is None and len(self.rec_vars_outputs) == 1:
      return list(self.rec_vars_outputs.values())[0]
    assert not self.rec_vars_outputs
    return None

  def post_process_final_rec_vars_outputs(self, rec_vars_outputs, seq_len):
    """
    :param dict[str,tf.Tensor] rec_vars_outputs:
    :param tf.Tensor seq_len: shape (batch,)
    :rtype: dict[str,tf.Tensor]
    """
    return rec_vars_outputs

  @classmethod
  def get_rec_initial_output(cls, batch_dim, name, output, rec_layer, initial_output=None, **kwargs):
    """
    If this layer is used inside a recurrent layer, this function specifies the
    output of frame t=-1, if it is needed.
    As arguments, we get the usual layer arguments.
    batch_dim is added because it might be special because of beam search.

    Note: This could maybe share code with :func:`RnnCellLayer.get_rec_initial_state`.

    :param tf.Tensor batch_dim: including beam size in beam search
    :param str name: layer name
    :param Data output: template
    :param returnn.tf.layers.rec.RecLayer rec_layer:
    :param str|float|int|tf.Tensor|None initial_output:
    :rtype: tf.Tensor
    """
    import numpy
    v = initial_output
    data = output
    if isinstance(v, tf.Tensor):
      return v
    if v is None and data.sparse:
      raise Exception(
        ("You must explicitly provide an initial output value for sparse data %r." % data) +
        (" E.g. '%s': {'initial_output': 'zeros'}." % name))
    if v is None:
      v = "zeros"
    bc_shape = [(d if (d is not None) else 1) for d in data.batch_shape]
    # Some other code might not support automatic broadcasting in the batch-axis. (Earlier example: concat_in_time)
    # Thus we will automatically unbroadcast here (at least the batch axis).
    # Note that there still might be other axes which we do not unbroadcast here.
    # Thus, concat_in_time was fixed now, and maybe we actually do not need this anymore.
    shape = list(bc_shape)
    shape[data.batch_dim_axis] = batch_dim
    if isinstance(v, (float, int)):
      with tf.name_scope("init_%s_const" % name):
        from returnn.tf.util.basic import constant_with_shape
        return tf.cast(constant_with_shape(v, shape=shape), dtype=data.dtype)
    if isinstance(v, LayerBase):
      v = v.output.copy_compatible_to(output)
      if output.beam:
        v = v.copy_extend_with_beam(output.beam)
      return v.placeholder
    assert isinstance(v, str)
    if v == "zeros":
      return tf.zeros(shape, dtype=data.dtype, name="init_%s_zeros" % name)
    elif v == "ones":
      return tf.ones(shape, dtype=data.dtype, name="init_%s_ones" % name)
    elif v == "var":
      assert not data.sparse
      assert numpy.prod(bc_shape) == data.dim
      with rec_layer.var_creation_scope():
        x = tf_compat.v1.get_variable(
          "init_%s_var" % name, shape=(data.dim,), dtype=data.dtype, initializer=tf.zeros_initializer())
      x = tf.reshape(x, bc_shape, name="init_%s_var_bc" % name)
      x = tf.tile(x, [batch_dim if (i == data.batch_dim_axis) else 1 for i in range(data.batch_ndim)],
                  name="init_%s_var_batch_bc" % name)
      return x
    elif v == "apply(0)":
      # We will apply the layer for the input 0 and use this as the initial state.
      # This code might be a bit unstable.
      kwargs = kwargs.copy()
      sources = kwargs.pop("sources")
      zeroed_sources = []
      for src in sources:
        assert isinstance(src, LayerBase)
        src_output = src.output.copy()
        if src_output.placeholder is not None:
          zeroed_src_shape = tf_util.get_shape(src_output.placeholder)
          zeroed_src_shape = [
            zeroed_src_shape[i] for i in range(src_output.batch_ndim)]  # type: typing.List[typing.Union[tf.Tensor,int]]
        else:
          zeroed_src_shape = [(d if (d is not None) else 1) for d in src_output.batch_shape]
        if src_output.batch_dim_axis is not None:
          zeroed_src_shape[src_output.batch_dim_axis] = batch_dim
        src_output.placeholder = tf.zeros(
          zeroed_src_shape, dtype=src_output.dtype,
          name="init_%s_zeros" % tf_util.get_valid_scope_name_from_str(src.name))
        src_output.sanity_check()
        zeroed_src = InternalLayer(name="%s_zeroed" % src.name, output=src_output, network=src.network)
        zeroed_sources.append(zeroed_src)
      layer = cls(name=name, output=output.copy(), sources=list(zeroed_sources), **kwargs)
      out = layer.output.placeholder
      out.set_shape(data.batch_shape)
      return out
    else:
      raise Exception("invalid initial output type %r for sub-layer %r" % (v, name))

  @classmethod
  def get_rec_initial_extra_outputs(cls, batch_dim, rec_layer, **kwargs):
    """
    :param tf.Tensor batch_dim: for this layer, might be with beam
    :param TFNetworkRecLayer.RecLayer rec_layer:
    :rtype: dict[str,tf.Tensor]
    """
    return {}

  @classmethod
  def get_rec_initial_extra_outputs_shape_invariants(cls, **kwargs):
    """
    :return: optional shapes for the tensors by get_rec_initial_extra_outputs
    :rtype: dict[str,tf.TensorShape]
    """
    return {}


class InternalLayer(LayerBase):
  """
  This is not supposed to be used by the user.
  It is used by some code to construct a wrapper layer or so.
  """


class DataNotAvailableLayer(InternalLayer):
  """
  This is a dummy layer that is created when the output template is flagged "not available for inference".
  The output template should be passed to the constructor to correctly forward the information
  in case any dependent output is exported with "register_as_extern_data".

  See :func:`returnn.tf.network._create_layer`
  """
  def __init__(self, layer_class, layer_desc, **kwargs):
    """
    :param type[LayerBase] layer_class:
    :param dict[str] layer_desc:
    """
    super(DataNotAvailableLayer, self).__init__(**kwargs)
    self.layer_class_ = layer_class
    self.layer_desc = layer_desc

  def get_sub_layer(self, layer_name):
    """
    :param str layer_name: name of the sub_layer (right part of '/' separated path)
    :rtype: LayerBase|None
    """
    from returnn.tf.network import TFNetwork
    cls = self.layer_class_
    assert issubclass(cls, LayerBase)
    res = cls.get_sub_layer_out_data_from_opts(layer_name=layer_name, parent_layer_kwargs=self.layer_desc)
    if not res:
      return None
    out, net, sub_layer_class = res
    assert isinstance(out, Data)
    assert isinstance(net, TFNetwork)
    assert issubclass(sub_layer_class, LayerBase)
    sub_layer_desc = {}  # unfortunately no way to know, this would be dynamic in general...
    return DataNotAvailableLayer(
      name="%s/%s" % (self.name, layer_name), network=net, output=out,
      layer_class=sub_layer_class, layer_desc=sub_layer_desc)


class WrappedInternalLayer(InternalLayer):
  """
  This is not supposed to be used by the user. Like :class:`InternalLayer`, only intended for internal usage.
  This layer is supposed to logically wrap another layer.
  """

  def __init__(self, base_layer, **kwargs):
    """
    :param LayerBase base_layer: the layer which we are wrapping
    """
    super(WrappedInternalLayer, self).__init__(**kwargs)
    self.base_layer = base_layer
    self.params.update(base_layer.params)  # maybe ReuseParams wants to access it or so

  def get_base_absolute_name_scope_prefix(self):
    """
    :rtype: str
    """
    return self.base_layer.get_base_absolute_name_scope_prefix()

  def get_absolute_name_scope_prefix(self):
    """
    :rtype: str
    """
    return self.base_layer.get_absolute_name_scope_prefix()

  def get_dep_layers(self):
    """
    :rtype: list[LayerBase]
    """
    return [self.base_layer]


class ReuseParams:
  """
  This is for parameter sharing, i.e. reusing existing `tf.Variable` objects in a new layer,
  instead of creating new variables.
  :func:`ReuseParams.from_config_dict` will be called via :func:`LayerBase.transform_config_dict`.
  """

  @classmethod
  def from_config_dict(cls, opts, network, get_layer):
    """
    This will be called via :func:`LayerBase.transform_config_dict` on the layer option `"reuse_params"`.

    :param str|dict[str]|None opts:
      If None, we will return None.
      If str, it will be interpret as a layer name.
      If dict, you can specify:
        "reuse_layer": layer name
        "map": dict where the keys are parameter names, and the values can be:
          A str would be interpret as a layer name.
          None would be interpret as the option `auto_create_missing`.
          A dict would specify :func:`ReuseParams.__init__` options.
            The option reuse_layer would be specified as a str, and represents a layer name.
    :param returnn.tf.network.TFNetwork network:
    :param ((str) -> LayerBase) get_layer: function to get or construct another layer
    :rtype: ReuseParams|None
    """
    if not opts:
      return None

    def optional_get_layer(layer_name):
      """
      :param str layer_name:
      :rtype: LayerBase|ReuseParams.LazyLayerResolver
      """
      from returnn.tf.network import NetworkConstructionDependencyLoopException, LayerNotFound
      try:
        return get_layer(layer_name)
      except (NetworkConstructionDependencyLoopException, LayerNotFound):
        # This dependency loop is not seen as critical. We allow it to be done later.
        # So any template construction of this layer should work.
        return ReuseParams.LazyLayerResolver(layer_name=layer_name, network=network, get_layer=get_layer)

    if isinstance(opts, str):  # share the whole layer
      return ReuseParams(reuse_layer=optional_get_layer(opts))
    assert isinstance(opts, dict)
    opts = opts.copy()
    if "reuse_layer" in opts:  # share the whole layer (+ possibly some more params)
      opts["reuse_layer"] = optional_get_layer(opts["reuse_layer"])
    if "map" in opts:  # share specific parameters
      assert isinstance(opts["map"], dict), "reuse_params['map'] should be a dict but is %s" % (type(opts["map"]),)
      opts["map"] = opts["map"].copy()
      for key, value in sorted(opts["map"].items()):
        if isinstance(value, str):
          value = {"reuse_layer": optional_get_layer(value)}
        elif value is None:
          value = {"auto_create_missing": True}
        else:
          assert isinstance(value, dict)
          value = value.copy()
          if "reuse_layer" in value:
            value["reuse_layer"] = optional_get_layer(value["reuse_layer"])
        opts["map"][key] = ReuseParams(**value)
    return ReuseParams(**opts)

  class LazyLayerResolver:
    """
    Unfortunately this is a bit tricky and difficult to do right.
    We want to support it because it can happen that e.g. in training, this is layer resolving is not needed,
    and then in search, it is needed, due to different dependencies.
    See :func:`test_reuse_params_map_custom_dep_loop` for an example.
    The params depend on a layer which is not constructed yet and cannot be constructed yet
    because of a dependency loop.
    Thus, here we again try to create it, and if we still get the dependency loop,
    we create the reused-params-layer based on dummy inputs, such that the variables/parameters get created
    and can be used now. Then, later, we are going to recreate the reused-params-layer.
    """

    def __init__(self, layer_name, network, get_layer):
      """
      :param str layer_name:
      :param returnn.tf.network.TFNetwork network:
      :param ((str) -> LayerBase) get_layer:
      """
      self.layer_name = layer_name
      self.network = network
      self.get_layer_func = get_layer
      self.var_scope = tf_compat.v1.get_variable_scope()

    def __repr__(self):
      return "<%s layer %r, net %r>" % (self.__class__.__name__, self.layer_name, self.network)

    def get_layer(self):
      """
      :rtype: LayerBase
      """
      from returnn.tf.network import NetworkConstructionDependencyLoopException, LayerNotFound
      from returnn.tf.util.basic import reuse_name_scope
      with reuse_name_scope(self.var_scope):
        try:
          return self.get_layer_func(self.layer_name)
        except (NetworkConstructionDependencyLoopException, LayerNotFound):
          return self.create_dummy_layer()

    def create_dummy_layer(self):
      """
      :rtype: LayerBase
      """
      from .basic import get_layer_class
      print(
        ("ReuseParams: layer %r does not exist yet and there is a dependency loop, " +
         "thus creating it on dummy inputs now") % self.layer_name,
        file=log.v4)

      layer_name = self.layer_name
      network = self.network
      with_time_dim = False
      while layer_name.startswith("base:") and network.parent_net:
        if network.parent_layer and network.parent_layer.output.have_time_axis():
          with_time_dim = True
        layer_name = layer_name[len("base:"):]
        network = network.parent_net

      # noinspection PyShadowingNames
      def get_dummy_input_layer(layer_name):
        """
        :param str layer_name:
        :rtype: LayerBase
        """
        if layer_name in network.layers:
          return network.layers[layer_name]
        output = None
        net = network

        # noinspection PyShadowingNames
        def opt_get_layer(layer_name):
          """
          :param str layer_name:
          :rtype: LayerBase
          """
          if layer_name in net.layers:
            return net.layers[layer_name]
          print("ReuseParams: non-existing layer %r in %r, ignoring..." % (layer_name, net), file=log.v4)
          return InternalLayer(
            name=layer_name, network=net,
            output=Data(
              name="LazyLayerResolver_dummy_output_%s" % layer_name,
              shape=(None, 1) if with_time_dim else ()))

        if self.network.parent_net is network and self.network.parent_layer:
          if layer_name.startswith(self.network.parent_layer.name + "/"):
            net = self.network
            layer_name = layer_name[len(net.parent_layer.name) + 1:]
            if layer_name in net.layers:
              # Don't return layer, could be inside loop and that wont work.
              output = net.layers[layer_name].output.copy_template()
              if not output.have_time_axis() and with_time_dim:
                output = output.copy_template_adding_time_dim()
        if not output:
          layer_desc_ = net.layers_desc[layer_name].copy()
          class_name_ = layer_desc_.pop("class")
          layer_class_ = get_layer_class(class_name_)
          layer_class_.transform_config_dict(layer_desc_, network=net, get_layer=opt_get_layer)
          # noinspection PyProtectedMember
          layer_desc_ = net._create_layer_layer_desc(name=layer_name, layer_desc=layer_desc_)
          output = layer_class_.get_out_data_from_opts(**layer_desc_).copy()
        output.beam = None
        output.placeholder = tf.zeros(
          [d or 1 for d in output.batch_shape], dtype=output.dtype, name="%s_dummy" % output.name)
        if not output.size_placeholder:
          output.size_placeholder = {}
        for i, dim in enumerate(output.shape):
          if dim is None and i not in output.size_placeholder:
            output.size_placeholder[i] = tf.ones([1], dtype=tf.int32, name="dummy_reuse_params_size")
        output.sanity_check()
        print("ReuseParams: creating dummy input %r with %r" % (layer_name, output), file=log.v4)
        return InternalLayer(name=layer_name, network=network, output=output)

      layer_desc = network.layers_desc[layer_name].copy()
      class_name = layer_desc.pop("class")
      layer_class = get_layer_class(class_name)
      layer_class.transform_config_dict(layer_desc, network=network, get_layer=get_dummy_input_layer)
      with reuse_name_scope(network.get_absolute_name_scope_prefix()[:-1], absolute=True):
        # noinspection PyProtectedMember
        return network._create_layer(
          name=layer_name, layer_class=layer_class, **layer_desc)

  # noinspection PyShadowingBuiltins
  def __init__(self, reuse_layer=None, map=None, custom=None, auto_create_missing=False):
    """
    :param LayerBase|ReuseParams.LazyLayerResolver|None reuse_layer:
    :param dict[str,ReuseParams]|None map:
    :param (**kwargs)->(tf.Tensor|tf.Variable) custom: see :func:`self.variable_custom_getter`
    :param bool auto_create_missing:
    """
    assert isinstance(reuse_layer, (LayerBase, ReuseParams.LazyLayerResolver)) or not reuse_layer
    self._reuse_layer = reuse_layer
    self.param_map = map
    self.custom_func = custom
    self.auto_create_missing = auto_create_missing

  def __repr__(self):
    return "<%s reuse_layer %r, map %r>" % (self.__class__.__name__, self._reuse_layer, self.param_map)

  @property
  def reuse_layer(self):
    """
    :rtype: LayerBase|None
    """
    if self._reuse_layer:
      if isinstance(self._reuse_layer, ReuseParams.LazyLayerResolver):
        self._reuse_layer = self._reuse_layer.get_layer()
      assert isinstance(self._reuse_layer, LayerBase)
      return self._reuse_layer
    return None

  def get_absolute_name_scope_prefix(self, base_layer, param):
    """
    :param LayerBase base_layer:
    :param tf.Variable param: e.g. "base_layer/rec/W"
    :return: e.g. "base_layer/" or "base_layer/rec/", always with "/" at end
    :rtype: str
    """
    # abs_scope_prefix: e.g. "current_layer/" or "current_layer/rec/"
    abs_scope_prefix = base_layer.get_absolute_name_scope_prefix()
    assert abs_scope_prefix.endswith("/")
    from returnn.tf.util.basic import get_current_var_scope_name
    # abs_cur_scope: e.g. "current_layer/rec/" or "current_layer/" or "current_layer/batch_norm/"
    abs_cur_scope = get_current_var_scope_name() + "/"
    assert abs_cur_scope.startswith(abs_scope_prefix)
    cur_scope_rel = abs_cur_scope[len(abs_scope_prefix):]  # e.g. "batch_norm/" or ""
    assert cur_scope_rel == "" or abs_cur_scope.endswith("/")
    assert param.name[-2:] == ":0"
    abs_param_name = param.name[:-2]
    abs_param_scope, _sep, param_name = abs_param_name.rpartition("/")
    assert param_name and (abs_param_scope + _sep).endswith(cur_scope_rel)
    if cur_scope_rel:
      param_name = cur_scope_rel + param_name
    del abs_param_scope
    assert abs_param_name.endswith("/" + param_name)
    if self.custom_func:  # Could be any base absolute name scope prefix, so just return what we have.
      return abs_param_name[:-len(param_name)]
    if self.param_map is not None and param_name in self.param_map:
      return self.param_map[param_name].get_absolute_name_scope_prefix(base_layer=base_layer, param=param)
    if self.reuse_layer and param_name in self.reuse_layer.params:
      reuse_layer_prefix = self.reuse_layer.get_absolute_name_scope_prefix()
      assert reuse_layer_prefix + param_name == abs_param_name
      return reuse_layer_prefix
    assert self.auto_create_missing
    base_layer_prefix = base_layer.get_absolute_name_scope_prefix()
    assert base_layer_prefix + param_name == abs_param_name
    return base_layer_prefix

  def get_variable_scope(self, base_layer, **kwargs):
    """
    :param LayerBase base_layer:
    :param kwargs: passed to tf.compat.v1.variable_scope
    :rtype: tf.compat.v1.VariableScope
    """
    def _variable_custom_getter(**kwargs_):
      return self.variable_custom_getter(base_layer=base_layer, **kwargs_)
    with tf_compat.v1.variable_scope(
          tf_compat.v1.get_variable_scope(), custom_getter=_variable_custom_getter, **kwargs) as scope:
      return scope

  def variable_custom_getter(self, getter, name, base_layer, **kwargs):
    """
    By TF docs, from :func:`_VariableStore.get_variable`:
    Callable that takes as a first argument the true getter,
    and allows overwriting the internal get_variable method.
    The signature of `custom_getter` should match that of this method,
    but the most future-proof version will allow for changes:
    `def custom_getter(getter, *args, **kwargs)`.  Direct access to
    all `get_variable` parameters is also allowed:
    `def custom_getter(getter, name, *args, **kwargs)`.  A simple identity
    custom getter that simply creates variables with modified names is:
    ```python
    def custom_getter(getter, name, *args, **kwargs):
      return getter(name + '_suffix', *args, **kwargs)
    ```
    In addition, we get the argument `base_scope_name`, via :func:`self.get_variable_scope`.

    :param (...)->tf.Variable getter:
    :param str name: absolute param name
    :param LayerBase base_layer: we expect that this is the prefix of ``name``
    :rtype: tf.Variable|tf.Tensor
    """
    abs_scope_prefix = base_layer.get_absolute_name_scope_prefix()
    assert not abs_scope_prefix or abs_scope_prefix.endswith("/")
    assert name.startswith(abs_scope_prefix)
    param_name = name[len(abs_scope_prefix):]  # e.g. "W" (not "rec/W")
    if self.custom_func:
      return self.custom_func(
        base_layer=base_layer, reuse_layer=self.reuse_layer, name=param_name, getter=getter, full_name=name, **kwargs)
    if self.param_map is not None:
      if not self.auto_create_missing:
        assert param_name in self.param_map
      if param_name in self.param_map:
        return self.param_map[param_name].variable_custom_getter(
          getter=getter, name=name, base_layer=base_layer, **kwargs)
    if self.reuse_layer:
      if not self.auto_create_missing:
        assert param_name in self.reuse_layer.params
      if param_name in self.reuse_layer.params:
        return self.reuse_layer.params[param_name]
    assert self.auto_create_missing
    return getter(name=name, **kwargs)


class SearchChoices(object):
  """
  In beam search, after expanding the beam and then selecting the N best (beam) (see :class:`ChoiceLayer`),
  when doing this multiple times, we need to keep reference where each beam came from,
  and what the current score is, etc.
  Also we could have multiple different such expansions & prunes via different :class:`ChoiceLayer`.
  This is what we keep track here.
  """

  def __init__(self, owner, beam_size, is_decided=False, keep_raw=False):
    """
    :param LayerBase owner:
    :param int beam_size:
    :param bool is_decided: by :class:`DecideLayer`
    :param bool keep_raw: by :class:`DecideKeepBeamLayer`
    """
    assert beam_size is not None
    self.owner = owner
    self._done_src_layer = False
    self._src_layer = None  # type: typing.Optional[LayerBase]
    self.src_beams = None  # type: typing.Optional[tf.Tensor]  # src beam index, (batch, beam)
    self.beam_size = beam_size
    self.beam_scores = None  # type: typing.Optional[tf.Tensor]  # (batch, beam)
    self.is_decided = is_decided
    self.keep_raw = keep_raw
    if not owner.output.beam:
      assert beam_size == 1
    else:
      assert owner.output.beam.beam_size == beam_size
      owner.network.register_search_choices_for_beam(beam=owner.output.beam, search_choices=self)

  def __repr__(self):
    def short(v):
      """
      :param LayerBase|tf.Tensor|None v:
      :return: short repr
      :rtype: str
      """
      if isinstance(v, LayerBase):
        return repr(v.name)
      if isinstance(v, tf.Tensor):
        if v.get_shape().ndims is not None:
          return "shaped:(%s)" % ",".join(map(str, v.get_shape().as_list()))
        return "unknown-ndim"
      return repr(v)
    s = " beam_size=%r" % self.beam_size
    if self._done_src_layer:
      s += " src_layer=%s" % short(self._src_layer)
    s += " beam_scores=%s" % short(self.beam_scores)
    if self.is_decided:
      s += " is_decided"
    if self.keep_raw:
      s += " keep_raw"
    return "<SearchChoices owner=%s%s>" % (short(self.owner), s)

  @property
  def src_layer(self):
    """
    :return: The layer where we had the last search choices.
    :rtype: LayerBase
    """
    if not self._done_src_layer:
      self._src_layer = self.owner.network.get_search_choices(base_search_choice=self.owner)
      self._done_src_layer = True
    return self._src_layer

  def set_beam_from_own_rec(self):
    """
    Assumes we have set self.owner, and uses those rec vars to set the beam scores.
    """
    self.set_beam_from_rec(self.owner.rec_vars_outputs)

  def set_beam_from_rec(self, rev_vars_outputs):
    """
    :param dict[str,tf.Tensor] rev_vars_outputs: e.g. via :class:`ChoiceLayer`
    """
    assert (
      rev_vars_outputs.get("choice_scores", None) is not None and
      rev_vars_outputs.get("choice_src_beams", None) is not None)
    self.beam_scores = rev_vars_outputs["choice_scores"]  # (batch, beam)
    self.src_beams = rev_vars_outputs["choice_src_beams"]  # (batch, beam)
    self.beam_scores.set_shape(self.src_beams.get_shape())

  def set_src_beams(self, src_beam_idxs):
    """
    :param tf.Tensor src_beam_idxs: source beam index, (batch, beam)
    """
    if isinstance(self.beam_size, int):
      src_beam_idxs.set_shape((None, self.beam_size))
    self.src_beams = src_beam_idxs
    self.owner.rec_vars_outputs["choice_src_beams"] = src_beam_idxs

  def set_beam_scores(self, scores):
    """
    :param tf.Tensor scores: (batch, beam) -> log score
     """
    if isinstance(self.beam_size, int):
      scores.set_shape((None, self.beam_size))
    self.beam_scores = scores
    self.owner.rec_vars_outputs["choice_scores"] = scores

  def get_src_choices_seq(self):
    """
    :return: all SearchChoices we depend on up to the root, including and starting with self
    :rtype: list[SearchChoices]
    """
    sources = [self]
    choice = self
    while True:
      src_layer = choice.src_layer
      if not src_layer:
        break
      assert isinstance(src_layer.search_choices, SearchChoices)
      choice = src_layer.search_choices
      if choice in sources:  # loop, can happen e.g. in RecLayer in the loop
        break
      sources.append(choice)
    return sources

  def get_beam_info(self):
    """
    :rtype: returnn.tf.util.data.SearchBeam|None
    """
    if self.owner.output.beam is None:
      assert self.beam_size == 1
      return None
    assert self.owner.output.beam.beam_size == self.beam_size
    return self.owner.output.beam

  def __eq__(self, other):
    return self is other

  def __ne__(self, other):
    return self is not other

  @staticmethod
  def compare(self, other):
    """
    Also see :func:`TFNetwork.get_search_choices.compare_layer`, which is basically the same.

    :param SearchChoices|None self:
    :param SearchChoices|None other:
    :return: 0 if equal, -1 if we are smaller, else 1
    :rtype: int
    """
    if self is other:
      return 0
    if self is None:
      return -1
    if other is None:
      return 1
    if self.keep_raw or other.keep_raw:
      return 0
    self_norm_layer = self.owner.get_normalized_layer()
    other_norm_layer = other.owner.get_normalized_layer()
    if self_norm_layer != self.owner and other_norm_layer != other.owner:
      assert self_norm_layer.search_choices and other_norm_layer.search_choices
      return SearchChoices.compare(self=self_norm_layer.search_choices, other=other_norm_layer.search_choices)
    self_src_choices = self.get_src_choices_seq()
    other_src_choices = other.get_src_choices_seq()
    if self in other_src_choices and other not in self_src_choices:
      return -1
    if other in self_src_choices and self not in other_src_choices:
      return 1
    from pprint import pformat
    raise Exception("Cannot compare search choices\n %r,\n %r\nwhich have traces:\n%s,\n%s" % (
      self, other, pformat(self_src_choices), pformat(other_src_choices)))

  def __cmp__(self, other):
    return self.compare(self, other)

  def __lt__(self, other):
    return self.__cmp__(other) < 0

  def __gt__(self, other):
    return self.__cmp__(other) > 0

  def translate_to_this_search_beam(self, sources):
    """
    :param LayerBase|list[LayerBase]|dict[str,LayerBase|object]|tuple[LayerBase|object]|T sources:
    :return: sources but all layers transformed when needed
    :rtype: T
    """
    from .basic import SelectSearchSourcesLayer
    d = sources
    if isinstance(d, dict):
      return {k: self.translate_to_this_search_beam(v) for (k, v) in d.items()}
    if isinstance(d, (tuple, list)):
      from returnn.util.basic import make_seq_of_type
      return make_seq_of_type(type(d), [self.translate_to_this_search_beam(v) for v in d])
    if isinstance(d, LayerBase):
      return SelectSearchSourcesLayer.select_if_needed(d, search_choices=self)
    return d

  @classmethod
  def translate_to_common_search_beam(cls, layer_desc):
    """
    :param list[LayerBase]|dict[str,LayerBase|object] layer_desc:
    :return: sources but all layers transformed when needed
    :rtype: list[LayerBase]|dict[str,LayerBase|object]
    """
    assert "_src_common_search_choices" not in layer_desc  # do not set this manually
    from tensorflow.python.util import nest
    layers_flat = [v for v in nest.flatten(layer_desc) if isinstance(v, LayerBase)]
    if len(layers_flat) <= 1:
      return layer_desc
    search_choicess = []
    for layer in layers_flat:
      if not layer.output.beam:
        continue
      search_choices = layer.get_search_choices()
      from pprint import pformat
      assert search_choices, "layer %r has beam %r but no search choices; from layer desc\n%s" % (
        layer, layer.output.beam, pformat(layer_desc))
      search_choicess.append(search_choices)
    if not search_choicess:
      return layer_desc
    from functools import cmp_to_key
    common_choices = max(search_choicess, key=cmp_to_key(cls.compare))
    layer_desc = layer_desc.copy()
    layer_desc["_src_common_search_choices"] = common_choices
    return common_choices.translate_to_this_search_beam(layer_desc)


class Loss(object):
  """
  Base class for all losses.
  """
  class_name = None  # type: str  # used by get_loss_class()
  recurrent = False  # if this is a frame-wise criteria, this will be False

  def __init__(self, base_network, use_flatten_frames=True,
               use_normalized_loss=False, custom_norm_factor=None,
               scale=1.0):
    """
    :param returnn.tf.network.TFNetwork base_network:
    :param bool use_flatten_frames: will use :func:`TFUtil.flatten_with_seq_len_mask`
    :param bool use_normalized_loss: the loss used in optimization will be normalized
    :param float|function|None custom_norm_factor:
    :param float scale: additional scale factor for the loss
    """
    self.base_network = base_network
    self.use_flatten_frames = use_flatten_frames
    self.layer = None  # type: typing.Optional[LayerBase]
    # All are initialized in self.init().
    self.output = None  # type: typing.Optional[Data]
    self.output_with_activation = None  # type: typing.Optional[OutputWithActivation]
    self.output_seq_lens = None  # type: typing.Optional[tf.Tensor]
    self.target = None  # type: typing.Optional[Data]
    self.target_seq_lens = None  # type: typing.Optional[tf.Tensor]
    self.output_flat = None  # type: typing.Optional[tf.Tensor]
    self.output_before_softmax_flat = None  # type: typing.Optional[tf.Tensor]
    self.target_flat = None  # type: typing.Optional[tf.Tensor]
    # Maybe make configurable. For now, same as in our Theano behavior.
    self.loss_norm_factor = None  # type: typing.Optional[tf.Tensor]
    self.use_normalized_loss = use_normalized_loss
    self.custom_norm_factor = custom_norm_factor
    self.scale = scale

  def __repr__(self):
    return "<%s %r>" % (self.__class__.__name__, self.layer or self.output)

  def _reduce_batch_time(self):
    """
    :return: In self.reduce_func, whether to expect that the loss is of shape (batch*time|time*batch,).
    :rtype: bool
    """
    if self.use_flatten_frames:
      return False  # we have used flatten_with_seq_len_mask. see self._flatten_or_merge
    if self.recurrent:
      return False
    if not self.output.have_time_axis():
      return False
    return True

  def _reduce_to_batch_time_with_mask(self, loss, normalize=False):
    """
    :param tf.Tensor loss: (batch*time,...) or (time*batch,...) depending if self.output is batch/time major
    :param bool normalize: for remaining dims. False -> use tf.reduce_sum, True -> use tf.reduce_mean
    :return: (batch*time,) or (time*batch,)
    :rtype: tf.Tensor
    """
    assert {self.output.batch_dim_axis, self.output.time_dim_axis} == {0, 1}
    if loss.get_shape().ndims > 1:
      reduce_func = tf.reduce_mean if normalize else tf.reduce_sum
      loss = reduce_func(loss, axis=list(range(1, loss.get_shape().ndims)))  # reduce remaining dims already
    mask = self.output.get_sequence_mask()  # e.g. (B,T)
    mask = tf.reshape(mask, [tf.shape(loss)[0]])
    loss = tf.where(mask, loss, tf.zeros_like(loss), "loss_masked")
    return loss

  def reduce_func(self, loss):
    """
    Reduces the frames. Currently the sum, and we do averaging later.
    We might change this logic at some point.
    Also, some code overwrites this function externally, e.g. with TFUtil.identity, to not do reducing.

    :param tf.Tensor loss: e.g. (batch*time,), or (time_flat,), or (batch*time,dim), etc
    :return: by default just a scalar. but this can be overwritten, to not reduce
    :rtype: tf.Tensor
    """
    if self._reduce_batch_time():
      # We expect to get (batch*time) or (time*batch) in the first dimension of the loss and the output.
      loss = self._reduce_to_batch_time_with_mask(loss)
    return tf.reduce_sum(loss)

  def reduce_to_batch(self, loss, normalize):
    """
    :param tf.Tensor loss: e.g. (batch*time,), or (time_flat,), or (batch*time,dim), etc
    :param bool normalize: reduce mean instead of reduce sum
    :return: (batch,)
    :rtype: tf.Tensor
    """
    if not self.recurrent and self.output.have_time_axis():
      assert not self.use_flatten_frames
      assert self._reduce_batch_time()
      # We expect to get (batch*time) or (time*batch) in the first dimension of the loss and the output.
      loss = self._reduce_to_batch_time_with_mask(loss, normalize=normalize)  # (batch*time,) or (time*batch,)
      loss.set_shape((None,))
      loss = tf.reshape(loss, tf.shape(self.output.placeholder)[:2])  # (batch,time) or (time,batch)
      loss = tf.reduce_sum(loss, axis=self.output.time_dim_axis)  # (batch,)
      if normalize:
        loss /= tf.cast(self.output.get_sequence_lengths(), tf.float32)
    else:
      # We expect that there is no time-dim, just a batch-dim. It could be like (batch,other_dims...) though.
      if loss.get_shape().ndims > 1:
        reduce_func = tf.reduce_mean if normalize else tf.reduce_sum
        loss = reduce_func(loss, axis=list(range(1, loss.get_shape().ndims)))  # (batch,)
    return loss

  @classmethod
  def transform_config_dict(cls, d, network, get_layer):
    """
    :param dict[str] d: will modify inplace, the loss_opts
    :param returnn.tf.network.TFNetwork network:
    :param ((str) -> LayerBase) get_layer: function to get or construct another layer

    Will modify `d` such that it becomes the kwargs for `self.__init__()`.
    Mostly leaves `d` as-is.
    This is used by `LayerBase.transform_config_dict`.
    """

  def init_by_layer(self, layer, layer_output_template=None):
    """
    :param LayerBase|None layer:
    :param Data|None layer_output_template: maybe alternative template
    """
    if layer_output_template and layer_output_template.have_time_axis() and not layer.output.have_time_axis():
      # It could be that the layer is from inside a RecLayer loop, and does not have a time dim.
      # In that case, use our template instead.
      layer_output = layer_output_template
    else:
      # Use the real layer.output instead.
      layer_output = layer.output
    if layer is self.layer and self.output is layer_output:
      return
    # noinspection PyProtectedMember
    self.init(
      output=layer_output,
      output_with_activation=layer.output_before_activation,
      target=layer._get_target_value(),
      layer=layer)

  def _flatten_or_merge(self, x, seq_lens, time_major):
    """
    :param tf.Tensor x: (B,T,...) or (T,B,...)
    :param tf.Tensor seq_lens: (B,)
    :param bool time_major:
    :return: (B*T|T*B|B',...)
    :rtype: tf.Tensor
    """
    if self.use_flatten_frames:
      from returnn.tf.util.basic import flatten_with_seq_len_mask
      return flatten_with_seq_len_mask(x, seq_lens, time_major=time_major)
    x_shape = tf.shape(x)
    x_shape = [x_shape[i] for i in range(x.get_shape().ndims)]
    if time_major != self.output.is_time_major:
      # We expect at various places (eg. reduce_to_batch) that the loss is the same as self.output.
      from returnn.tf.util.basic import swapaxes
      x = swapaxes(x, 0, 1)  # (B,T,...) or (T,B,...)
    return tf.reshape(x, [x_shape[0] * x_shape[1]] + x_shape[2:], name="merge_batch_time")

  def init(self, output, output_with_activation=None, target=None, layer=None):
    """
    :param Data output: generated output
    :param OutputWithActivation|None output_with_activation:
    :param Data target: reference target from dataset
    :param LayerBase|None layer:
    """
    flatten_or_merge = self._flatten_or_merge
    with tf.name_scope("loss_init"):
      self.layer = layer
      if target:
        if output.beam:
          if target.beam != output.beam:
            target = target.copy_extend_with_beam(output.beam)
        else:
          assert not target.beam
      if output.feature_dim_axis is not None and output.feature_dim_axis != output.batch_ndim - 1:
        if output_with_activation:
          from returnn.tf.util.basic import move_axis
          output_with_activation = OutputWithActivation(x=move_axis(output_with_activation.x,
                                                                    output.feature_dim_axis, -1),
                                                        act_func=output_with_activation.act_func)
        output = output.copy_with_feature_dim_axis(-1)
      self.output = output
      self.output_with_activation = output_with_activation
      self.target = target
      # Flat variants are with batch,time collapsed into one, masked via seq_lens.
      self.output_flat = None
      self.output_before_softmax_flat = None
      self.target_flat = None
      self.output_seq_lens = None
      self.target_seq_lens = None
      self.loss_norm_factor = 1.0
      if self.output.have_time_axis():
        self.output_seq_lens = output.get_sequence_lengths()
        time_and_batch_dims = (self.output.time_dim_axis, self.output.batch_dim_axis)
        assert time_and_batch_dims in [(0, 1), (1, 0)], (
          "output time-batch-dim unexpected: %r (target %r)" % (self.output, self.target))
        if output_with_activation and output_with_activation.act_func is tf.nn.softmax:
          self.output_before_softmax_flat = flatten_or_merge(
            output_with_activation.x, self.output_seq_lens, time_major=output.is_time_major)
        else:
          self.output_flat = flatten_or_merge(output.placeholder, self.output_seq_lens, time_major=output.is_time_major)
          self.output_flat.set_shape(tf.TensorShape(output.shape))
        if target:
          assert target.have_time_axis()
          self.target_seq_lens = target.get_sequence_lengths()
          self.target_flat = flatten_or_merge(target.placeholder, self.target_seq_lens, time_major=target.is_time_major)
          self.loss_norm_factor = 1.0 / tf.cast(tf.reduce_sum(self.target_seq_lens), tf.float32)
        else:
          self.loss_norm_factor = 1.0 / tf.cast(tf.reduce_sum(self.output_seq_lens), tf.float32)
      else:  # no time axis
        if output_with_activation and output_with_activation.act_func is tf.nn.softmax:
          self.output_before_softmax_flat = output_with_activation.x
        else:
          self.output_flat = output.placeholder
        if self.output.have_batch_axis():
          self.loss_norm_factor = (
            1.0 / tf.cast(tf.shape(self.output.placeholder)[self.output.batch_dim_axis], tf.float32))
        else:
          self.loss_norm_factor = 1.0
        if target:
          assert not self.target.have_time_axis()
          self.target_flat = target.placeholder
      if self.custom_norm_factor is not None:
        if callable(self.custom_norm_factor):
          self.loss_norm_factor = self.custom_norm_factor(self=self, output=output, layer=layer)
        else:
          assert isinstance(self.custom_norm_factor, float)
          self.loss_norm_factor = self.custom_norm_factor
      self._check_init()

  def _check_init(self):
    """
    Does some checks on self.target and self.output, e.g. if the dense shapes matches.
    You can overwrite this if those checks don't make sense for your derived loss class.
    """
    if not self.target:
      return
    assert self.target.ndim_dense == self.output.ndim_dense, (
      "Number of dimensions mismatch. Target: %s, output: %s" % (self.target, self.output))
    expected_output_dim = self.get_auto_output_layer_dim(self.target.dim)
    assert expected_output_dim == self.output.dim, (
      "Expected output dim is %r but the output has dim %r. " % (expected_output_dim, self.output.dim) +
      "Target: %s, output: %s" % (self.target, self.output))

  def get_error(self):
    """
    :return: frame error rate as a scalar value with the default self.reduce_func (see also self.get_value)
    :rtype: tf.Tensor
    """
    with tf.name_scope("loss_frame_error"):
      assert self.output.ndim_dense == self.target.ndim_dense
      from returnn.tf.util.basic import check_input_ndim, check_shape_equal
      output_flat = self.output_before_softmax_flat
      if output_flat is None:
        output_flat = self.output_flat
      output_flat = check_input_ndim(output_flat, ndim=2)
      last_dim = tf.rank(output_flat) - 1  # should be 1
      if self.target.sparse:
        target_label = check_input_ndim(self.target_flat, ndim=1)
      else:
        target_flat = check_shape_equal(self.target_flat, output_flat)
        target_label = tf.cast(tf.argmax(target_flat, axis=last_dim), tf.int32)
      output_label = tf.cast(tf.argmax(output_flat, axis=last_dim), target_label.dtype)
      not_equal = tf.not_equal(output_label, target_label)
      return self.reduce_func(tf.cast(not_equal, tf.float32))

  def get_value(self):
    """
    :return: self.reduce_func(loss), which is usually a scalar with the default as if does tf.reduce_sum.
      float32 value. it should *not* be normalized over frames,
      as this will be calculated in :func:`TFEngine.Runner._collect_eval_info`.
    :rtype: tf.Tensor|None
    """
    raise NotImplementedError

  def get_normalization_factor(self):
    """
    :return: factor as a float scalar, usually 1.0 / num_frames. see self.reduce_func.
    :rtype: tf.Tensor
    """
    assert self.loss_norm_factor is not None, "init not called?"
    return tf.convert_to_tensor(self.loss_norm_factor)

  @classmethod
  def get_auto_output_layer_dim(cls, target_dim):
    """
    :param int target_dim:
    :return: normally just the same as target_dim. e.g. for CTC, we would add 1 for the blank label
    :rtype: int
    """
    return target_dim

  @classmethod
  def get_default_target(cls, extern_data):
    """
    :param TFNetwork.ExternData extern_data:
    :return: default target name, or None if this loss does not have a target
    :rtype: str|None
    """
    return extern_data.default_target
