
"""
Provides :class:`Data`, :class:`Dim`, :class:`SearchBeam`.

See :ref:`data` for some higher-level description.
"""

from __future__ import print_function, division

import os
import typing
import tensorflow as tf
import traceback

from returnn.util.basic import NotSpecified, Entity
import returnn.tf.compat as tf_compat


class Dim(object):
  """
  This identifies one axis/dimension, like a time-dimension, etc.
  This was called ``DimensionTag`` earlier, and referred to as dimension tag.

  This is used by :class:`Data`. See :func:`Data.dim_tags`.
  This would be passed as ``dim_tags`` when creating a :class:`Data` instance.

  It is not to specify the specific axis in a specific Data/tensor,
  but to specify the content and dimension.
  I.e. if we have the same Dim for two Data instances,
  the dimensions should match. I.e.:

      data1.get_dim_tag(i) == data2.get_dim_tag(j)
        =>  tf.shape(data1.placeholder)[i] == tf.shape(data2.placeholder)[j]

  This also includes further information such as sequence lengths
  or a vocabulary.

  We differentiate between the batch dim, spatial dim or feature dim,
  although that is just flag and in many contexts there is no real difference
  between a spatial dim and a feature dim (the batch dim is often handled differently).

  See :func:`SpatialDim` and :func:`FeatureDim` as easy wrappers to create dim tags for the user.
  """

  class Types:
    """
    Defines possible values for ``kind``.
    """
    Unspecified = None
    Batch = Entity("batch")
    Spatial = Entity("spatial")  # also time
    Time = Spatial  # we don't treat this as different
    Feature = Entity("feature")
    Types = (Batch, Spatial, Feature)

  _creation_counter = 0

  def __init__(self, kind=Types.Unspecified, description=None,
               dimension=None,
               vocab=None,
               dyn_size=None, dyn_size_ext=None,
               undefined=False, generic=False, special=False,
               auto_generated=False,
               match_priority=0,
               derived_from_tag=None, derived_from_op=None,
               batch=None, control_flow_ctx=None,
               src_data=None, src_axis=None):
    """
    :param Entity|None kind:
    :param str|None description: the description should be unique
    :param int|None dimension:
    :param returnn.datasets.util.vocabulary.Vocabulary|None vocab:
    :param tf.Tensor|None dyn_size: e.g. seq_len, (batch,)
    :param Data|None dyn_size_ext: seq_len or extended
    :param bool undefined: When this is specified as `None` by the user via `shape`.
    :param bool generic: This can not be a dim tag of :class:`Data` but it would match to other existing dim tags
      :func:`Data.get_axis_from_description` and :func:`Dim.is_equal`.
    :param bool special: Like `generic`, this can not be a dim tag of :class:`Data`.
      But this dim tag also does not match anything except itself.
      So it can be used to represent special placeholders with special meanings like ``single_step``.
    :param bool auto_generated: This is auto-generated by RETURNN because it was not explicitly specified by the user.
      E.g. for ConvLayer and others. This implies certain behavior on equality, such as comparing the description,
      to allow for several independent creations of the dim tag during template construction.
    :param Dim|None derived_from_tag:
      Whether this new tag is reduced, down/up sampled, padded etc from this given other tag.
      In situations where dim tags are being matched (Data.get_common_data),
      the behavior is to consider them as equal,
      and assume that the chain of operations (e.g. padding + valid conv) results in the same dim.
    :param Dim.Op|None derived_from_op:
    :param int match_priority: when there is ambiguity between multiple dim tags, this value defines the order
      in which the dimension are assigned to their matching counterparts.
      A dimension tag with a higher priority value is assigned first.
      E.g. for a square matrix used for a linear transformation, the reduce dim tag should have a higher priority.
    :param BatchInfo|None batch: for batch-dim, or dynamic dims per batch
    :param ControlFlowContext|None control_flow_ctx:
    :param Data|None src_data:
    :param int|None src_axis:
    """
    assert kind is None or (isinstance(kind, Entity) and kind in self.Types.Types)
    assert dimension is None or isinstance(dimension, int)
    assert description is None or isinstance(description, str)
    self.kind = kind
    self.description = description
    self.dimension = dimension
    self._vocab = vocab
    self.same_as = None  # type: typing.Optional[Dim]
    self._same_as_tb = None  # type: typing.Optional[traceback.StackSummary]  # for debugging
    self.derived_from_tag = derived_from_tag
    self.derived_from_op = derived_from_op
    if derived_from_op and not derived_from_op.output:
      derived_from_op.output = self
    self.match_priority = match_priority
    if src_data:
      assert isinstance(src_data, Data) and isinstance(src_axis, int)
    if not batch and dyn_size_ext:
      batch = dyn_size_ext.batch
    self.batch = batch
    self.control_flow_ctx = control_flow_ctx
    self.src_data = src_data
    self.src_axis = src_axis
    if dyn_size_ext and not dyn_size_ext.batch and batch:
      dyn_size_ext.batch = batch
    if dyn_size_ext:
      assert batch == dyn_size_ext.batch
    self.dyn_size_ext = dyn_size_ext  # type: typing.Optional[Data]
    self._dyn_size_same = set()  # set of TensorRef
    self._undefined = undefined
    self.generic = generic
    self.special = special
    self.auto_generated = auto_generated
    # We can have different tag variants per batch info (e.g. with beam), or per control flow ctx.
    # They each have same_as = self. The same_base should have the base (global) batch info.
    self._same_for_batch_ctx = {}  # type: typing.Dict[typing.Tuple[BatchInfo,typing.Optional[ControlFlowContext]],Dim]  # nopep8
    if dyn_size is not None:
      assert not dyn_size_ext
      self.dyn_size = dyn_size
    self._creation_idx = Dim._creation_counter
    Dim._creation_counter += 1

  def __repr__(self):
    return "Dim{%s}" % self.short_repr()

  def short_repr(self):
    """
    :return: some short repr
    :rtype: str
    """
    if self.is_batch_dim():
      return "B"  # Data.__repr__ will additionally give info on the batch
    desc = "%s%r" % ("F" if self.is_feature_dim() else "", self.get_same_base().description)
    if self.special:
      desc += "!"
    elif self.dimension is not None:
      desc += "(%i%s)" % (self.dimension, "*" if self.generic else "")
    else:
      if self.dyn_size_ext:
        desc += "[%s%s]" % (
          ",".join(self.dyn_size_ext.get_batch_axes_short_description(special_axes=False)),
          "*" if self.generic else "")
      else:
        desc += "[*]" if self.generic else "[?]"
      if self.control_flow_ctx:
        desc += "{ctx=%s}" % self.control_flow_ctx.repr_inner()
    return desc

  def __getstate__(self):
    d = vars(self).copy()
    d["batch"] = None
    d["_same_as_tb"] = None
    d["_same_for_batch_ctx"] = {}
    d["dyn_size_ext"] = None
    d["kind"] = self.kind.name if self.kind else None
    return d

  def __setstate__(self, state):
    self.__dict__.update(state)
    if self.kind is not None:
      self.kind = {v.name: v for v in Dim.Types.Types}[self.kind]
    if self.is_batch_dim():
      self.same_as = batch_dim

  def __copy__(self):
    """
    Normally we would not want to get a new tag with ``tag != copy(tag)``.
    https://github.com/rwth-i6/returnn/issues/860

    See :func:`Dim.copy` if you explicitly want a copy.

    :return: self
    :rtype: Dim
    """
    return self

  def __deepcopy__(self, memo=None):
    """
    Normally we would not want to get a new tag with ``tag != deepcopy(tag)``.
    https://github.com/rwth-i6/returnn/issues/860

    See :func:`Dim.copy` if you explicitly want a copy.

    :param memo:
    :return: self
    :rtype: Dim
    """
    return self

  def copy(self, same_as_self=True, description=None, kind=None, match_priority=None):
    """
    :param bool same_as_self:
    :param str|None description: new description
    :param Entity|None kind: if set, overwrites self.kind
    :param int|None match_priority:
    :return: copy, maybe as new kind. setting same_as to self
    :rtype: Dim
    """
    assert self.can_be_used_as_dim()
    if not same_as_self:
      assert description is not None, "%s copy with not same_as_self should have a new description" % self
    tag = Dim(
      kind=kind or self.kind, description=description or self.description,
      match_priority=match_priority if match_priority is not None else self.match_priority,
      dimension=self.dimension, dyn_size_ext=self.dyn_size_ext,
      batch=self.batch,
      src_data=self.src_data, src_axis=self.src_axis)
    if same_as_self:
      tag.same_as = self  # not declare_same_as, none of the extra checks needed
      tag._same_as_tb = traceback.extract_stack()
    return tag

  def _can_use_in_ctx(self, ctx):
    """
    :param ControlFlowContext|None ctx:
    :rtype: bool
    """
    if self.control_flow_ctx == ctx:
      return True
    if not ControlFlowContext.is_parent_or_same(self.control_flow_ctx, ctx):
      return False
    assert ctx
    # E.g. ctx == loop(time_dim), when self.control_flow_ctx == None,
    # we can use self in ctx, iff time_dim not in self.dyn_size_ext.dim_tags.
    # We can only do this check if we know about dyn_size_ext.
    if not self.dyn_size_ext:
      return False
    parent_dims = ControlFlowContext.collect_parent_dims(ctx)
    for dim in self.dyn_size_ext.dim_tags:
      if dim in parent_dims:
        return False
    return True

  def _validate_in_current_graph(self):
    """
    :rtype: bool
    """
    tensor = None
    if self.batch:
      batch_base = self.batch.get_global_base()
      if batch_base.is_global_batch():
        tensor = batch_base.get_global_batch_dim().size
    if not isinstance(tensor, tf.Tensor):
      if self.dyn_size_ext and self.dyn_size_ext.placeholder is not None:
        tensor = self.dyn_size_ext.placeholder
    if isinstance(tensor, tf.Tensor) and not tf_compat.executing_eagerly():
      from returnn.tf.util import basic as tf_util
      g = tf_util.get_root_graph()
      if tf_util.get_root_graph(tensor.graph) is not g:  # maybe from an earlier run which reuses the dim tag
        # Reset and cleanup.
        if self.dyn_size_ext:
          self.dyn_size_ext = self.dyn_size_ext.copy_template()
          self.dyn_size_ext.batch = None
        same_base = self.get_same_base()
        same_base._same_for_batch_ctx.pop((self.batch, self.control_flow_ctx), None)
        self.batch = None  # it is invalid in the new graph
        self.control_flow_ctx = None  # also invalid
        return False
    return True

  def _maybe_update(self):
    if self.is_batch_dim():
      return
    if isinstance(self.dimension, int):
      return
    if not self.batch:
      if self.dyn_size_ext and self.dyn_size_ext.batch:
        self.batch = self.dyn_size_ext.batch
      else:
        return
    same_base = self.get_same_base()
    key = (self.batch, self.control_flow_ctx)
    if self.dyn_size_ext and key not in same_base._same_for_batch_ctx:
      same_base._same_for_batch_ctx[key] = self
    # Check if we can find more
    if key in same_base._same_for_batch_ctx:
      same = same_base._same_for_batch_ctx[key]
      if same is not self:
        if same.dyn_size_ext and not self.dyn_size_ext:
          self.dyn_size_ext = same.dyn_size_ext
        if same.dyn_size_ext and same.dyn_size_ext.placeholder is not None:
          if self.dyn_size_ext.placeholder is None:
            self.dyn_size_ext = same.dyn_size_ext
        if self.dyn_size_ext and not same.dyn_size_ext:
          same.dyn_size_ext = self.dyn_size_ext
        if self.dyn_size_ext and self.dyn_size_ext.placeholder is not None:
          if not same.dyn_size_ext or same.dyn_size_ext.placeholder is None:
            same.dyn_size_ext = self.dyn_size_ext

  def get_for_batch_ctx(self, batch, ctx, allow_none=False):
    """
    :param BatchInfo batch:
    :param ControlFlowContext|None ctx:
    :param bool allow_none:
    :rtype: Dim|None
    """
    assert self.can_be_used_as_dim()
    if self.batch == batch and self.control_flow_ctx == ctx and self.dyn_size_ext:
      self._validate_in_current_graph()
      self._maybe_update()
      if self.batch == batch and self.control_flow_ctx == ctx and self.dyn_size_ext:  # check again
        return self
    if self.is_batch_dim():
      # We ignore the ctx for the batch dim currently.
      if self.batch == batch:
        return self
      return Dim(kind=Dim.Types.Batch, description="batch:%s" % batch.short_repr(), batch=batch)
    if self.dimension is not None:
      # If static dim, no effect.
      assert not self.batch
      return self
    if batch.is_broadcast():
      return self  # just leave as-is. should not matter.
    same_base = self.get_same_base()
    same_base._validate_in_current_graph()
    # Might be uninitialized in some cases. Assume batch is global.
    if not same_base.batch:
      batch_base = batch.get_global_base()
      if same_base.dyn_size_ext:
        assert batch == batch_base
        same_base.batch = batch
        assert not same_base.dyn_size_ext.batch or same_base.dyn_size_ext.batch == batch
        same_base.dyn_size_ext.batch = batch
      else:
        same_base.batch = batch_base
    if same_base.dyn_size_ext:
      assert same_base.batch == same_base.dyn_size_ext.batch
      assert same_base.control_flow_ctx == same_base.dyn_size_ext.control_flow_ctx
    for ctx_ in ControlFlowContext.abs_ctx_stack_with_root(ctx):
      tag = same_base._same_for_batch_ctx.get((batch, ctx_), None)
      if tag and tag._can_use_in_ctx(ctx) and tag._validate_in_current_graph():
        return tag
    if same_base.batch == batch and same_base._can_use_in_ctx(ctx) and same_base.dyn_size_ext:
      return same_base
    # Ok, nothing matching found.
    dyn_size_ext = None
    # Maybe we have sth with the base batch without beam or padded batch which we can extend.
    if batch != batch.get_global_base():
      batch_base = batch.get_global_base()
      base_can_use_in_ctx = None
      if same_base.batch == batch_base and same_base._can_use_in_ctx(ctx) and same_base.dyn_size_ext:
        base_can_use_in_ctx = same_base
      else:
        for ctx_ in ControlFlowContext.abs_ctx_stack_with_root(ctx):
          tag = same_base._same_for_batch_ctx.get((batch_base, ctx_), None)
          if tag and tag._can_use_in_ctx(ctx) and tag._validate_in_current_graph() and tag.dyn_size_ext:
            base_can_use_in_ctx = tag
            break
      if base_can_use_in_ctx and base_can_use_in_ctx.dyn_size_ext:
        if base_can_use_in_ctx.dyn_size_ext.have_batch_axis():
          # The same_base has some dyn size without any beam nor control flow context.
          # We can expand it to the current beam, or extend by padded batch.
          dyn_size_ext = base_can_use_in_ctx.dyn_size_ext.copy_extend_batch(batch)
          if batch.beam:
            dyn_size_ext = base_can_use_in_ctx.dyn_size_ext.copy_extend_with_beam(batch.beam)
          assert dyn_size_ext.batch == batch
          if dyn_size_ext.placeholder is not None:
            beam_expanded_base_data = getattr(dyn_size_ext.placeholder, "_RETURNN_beam_expanded_base_data", None)
            if batch.beam:
              assert beam_expanded_base_data
            # Note: The beam expansion used tiling, which can be cached.
            # This means that we could end up with the same size tensor (placeholder) for multiple different beams,
            # when there are different beams with same beam size!
            # This breaks the current logic in get_tag_from_size_tensor.
            # As a workaround, we make an explicit new tensor here.
            from .basic import get_valid_scope_name_from_str, same_control_flow_ctx
            with same_control_flow_ctx(dyn_size_ext.placeholder):
              dyn_size_ext.placeholder = tf.identity(
                dyn_size_ext.placeholder,
                name=get_valid_scope_name_from_str("%s_get_for_batch_ctx_%s" % (dyn_size_ext.name, batch.short_repr())))
            if batch.beam:
              dyn_size_ext.placeholder._RETURNN_dyn_size_beam = batch.beam
              dyn_size_ext.placeholder._RETURNN_beam_expanded_base_data = beam_expanded_base_data
    if not dyn_size_ext and allow_none:
      return None
    dim_tag = Dim(
      kind=self.kind, description=self.description, dimension=self.dimension,
      auto_generated=self.auto_generated,
      batch=batch, control_flow_ctx=dyn_size_ext.control_flow_ctx if dyn_size_ext else ctx,
      dyn_size_ext=dyn_size_ext)
    dim_tag.same_as = same_base
    dim_tag._same_as_tb = traceback.extract_stack()
    if dyn_size_ext and dyn_size_ext.placeholder is not None:
      if Dim.get_tag_from_size_tensor(dyn_size_ext.placeholder) is None:
        dim_tag.set_tag_on_size_tensor(dyn_size_ext.placeholder, batch=batch)
    same_base._same_for_batch_ctx[(dim_tag.batch, dim_tag.control_flow_ctx)] = dim_tag
    return dim_tag

  def set_dyn_size_ext_for_batch_ctx(self, batch, ctx, dyn_size_ext):
    """
    :param BatchInfo batch:
    :param ControlFlowContext|None ctx:
    :param Data dyn_size_ext:
    """
    assert self.can_be_used_as_dim()
    same = self.get_for_batch_ctx(batch, ctx)
    same.dyn_size_ext = dyn_size_ext
    self._maybe_update()

  def get_dyn_size_ext_for_batch_ctx(self, batch, ctx):
    """
    :param BatchInfo|None batch:
    :param ControlFlowContext|None ctx:
    :rtype: Data|None
    """
    assert self.can_be_used_as_dim()
    if not batch and self.batch:
      # Assume global batch.
      batch = self.batch.get_global_base()
    if not batch:
      # This is usually not valid. However, this case can happen early at initialization.
      assert batch == self.batch and ctx == self.control_flow_ctx
      return self.dyn_size_ext
    same = self.get_for_batch_ctx(batch, ctx, allow_none=True)
    if not same:
      return None
    return same.dyn_size_ext

  @property
  def dyn_size(self):
    """
    :return: dyn size / seq len (usually of shape [B]), or None
      If the dyn size can potentially be of a different shape, directly access dyn_size_ext.
    :rtype: tf.Tensor|None
    """
    if self.dimension is None and (not self.dyn_size_ext or self.dyn_size_ext.placeholder is None):
      # Try to complete.
      self.complete_dyn_size()
    if self.dyn_size_ext:
      return self.dyn_size_ext.placeholder
    return None

  @dyn_size.setter
  def dyn_size(self, dyn_size):
    """
    Also see :func:`set_dyn_size_ext_for_batch_ctx`.

    :param tf.Tensor dyn_size:
    """
    assert self.can_be_used_as_dim()
    assert isinstance(dyn_size, tf.Tensor) and dyn_size.shape.ndims == 1
    if self.dyn_size_ext:
      if self.dyn_size_ext.placeholder is None:
        self.dyn_size_ext.placeholder = dyn_size
        return
      # Do not allow resetting it to sth different.
      assert self.dyn_size_ext.placeholder is dyn_size
      return
    beam = getattr(dyn_size, "_RETURNN_dyn_size_beam", None)
    self.dyn_size_ext = Data(
      name=("%s:dyn_size" % self.description) if self.description else dyn_size.op.name,
      dtype=Data.size_dtype, placeholder=dyn_size, shape=(), batch_dim_axis=0,
      batch=self.batch, beam=beam, control_flow_ctx=self.control_flow_ctx)
    other = Dim.get_tag_from_size_tensor(dyn_size)
    if other:
      self.declare_same_as(other)
    else:
      self.set_tag_on_size_tensor(dyn_size)

  def is_batch_dim(self):
    """
    :return: whether this dim tag is of kind batch
    :rtype: bool
    """
    return self.kind == Dim.Types.Batch

  def is_feature_dim(self):
    """
    :return: whether this dim tag is of kind feature
    :rtype: bool
    """
    return self.kind == Dim.Types.Feature

  def is_spatial_dim(self):
    """
    :return: whether this dim tag is of kind spatial
    :rtype: bool
    """
    return self.kind == Dim.Types.Spatial

  def is_dim_known(self):
    """
    :return: whether we know the dimension; basically whether this is defined
      (although `not self.undefined` is defined slightly differently)
    :rtype: bool
    """
    if self.is_batch_dim():
      return True
    if self.dimension is not None:
      return True
    if self.dyn_size_ext:
      return True
    base = self.get_same_base()
    for _, other in base._same_for_batch_ctx.items():
      if other.dyn_size_ext:
        return True
    return False

  def is_dynamic(self):
    """
    :return: whether the dim is not static. usually means that it has seq lengths
    :rtype: bool
    """
    return self.dimension is not None

  def can_be_used_as_dim(self):
    """
    :return: whether this can be used as a dim in :class:`Data`, i.e. it is not generic or special
    :rtype: bool
    """
    return not self.generic and not self.special

  def is_same_size_tensor(self, x):
    """
    :param tf.Tensor x:
    :return: whether this dim tag for this specific batch (incl beam) is the same as the given size
    :rtype: bool
    """
    if x is self.dyn_size:
      return True
    from .basic import TensorRef
    if TensorRef(x) in self._dyn_size_same:
      return True
    return False

  def set_tag_on_size_tensor(self, x, batch=None, same_as_before=False):
    """
    This function is used
    to couple a tf.Tensor instance representing the dyn size
    with the dim tag.

    This is usually a newly created dim tag,
    which is yet unset.

    It is also used to couple an existing dim tag with other dyn sizes
    which just differ by an expansion of the batch (e.g. search beam).

    See also :func:`get_tag_from_size_tensor`.
    Also see :func:`set_dyn_size_ext_for_batch_ctx`.

    :param tf.Tensor x:
    :param BatchInfo|None batch:
    :param bool same_as_before: implies it was set before, and the new size is the same.
      e.g. it could be some identity with added checks, or other change.
    :return: self or new dim tag
    :rtype: Dim
    """
    from .basic import TensorRef
    assert self.can_be_used_as_dim()
    # It's unusual if self.dimension is not None, but let's accept that.
    if hasattr(x, "_is_size_of_dim_tag"):
      # noinspection PyProtectedMember
      assert x._is_size_of_dim_tag in (None, self)
    # If we already have another dyn size set or different batch, create a new Dim instance.
    if self.batch and batch and self.batch != batch:
      assert not same_as_before  # it cannot be the same when it is another batch...
      new_dim_tag = self.get_for_batch_ctx(batch=batch, ctx=self.control_flow_ctx)
      new_dim_tag.set_tag_on_size_tensor(x, batch=batch)
      return new_dim_tag
    if self.dyn_size is not None and self.dyn_size is not x:
      if TensorRef(x) in self._dyn_size_same:
        pass  # ok, pass on
      elif same_as_before:
        self._dyn_size_same.add(TensorRef(x))
        # And now pass on.
      else:
        assert self.batch and batch
        # It's not clear what to do. We could create a new dim tag, but the sizes might be different.
        # Usually we should not get here.
        # So for now, just error.
        from .basic import format_graph_output
        raise Exception("\n".join([
          "%r (%r) already has size %r, and another incompatible size %r (batch %r) is being assigned." % (
            self, self.description, self.dyn_size, x, batch),
          "\nNew size computation graph:",
          format_graph_output(x, max_depth=3),
          "\nThis is maybe the result of an incorrect declare_same_as. Traceback of declare_same_as:",
          "".join(self._same_as_tb.format()) if self._same_as_tb else ("same_as = %s" % self.same_as)]))
    if batch and getattr(x, "_RETURNN_dyn_size_beam", None):
      assert batch.beam == getattr(x, "_RETURNN_dyn_size_beam")
    if self.batch and batch:
      assert self.batch == batch
    elif batch and not self.batch:
      self.batch = batch  # overtake
    if getattr(x, "_is_size_of_dim_tag", None) is None:
      setattr(x, "_is_size_of_dim_tag", self)
    if not self.is_batch_dim() and self.dyn_size is None:
      self.dyn_size = x
    return self

  @classmethod
  def get_tag_from_size_tensor(cls, x):
    """
    :param tf.Tensor x: size tensor. has been set before via :func:`set_tag_on_size_tensor`
    :rtype: Dim|None
    """
    return getattr(x, "_is_size_of_dim_tag", None)

  def complete_dyn_size(self):
    """
    In case we can calculate the dyn size, do that now.
    """
    if self.dimension is not None:
      return
    if self.dyn_size_ext and self.dyn_size_ext.placeholder is not None:
      return
    same_base = self.get_same_base()
    op = self.derived_from_op or same_base.derived_from_op
    if op:
      from returnn.tf.util import basic as tf_util
      kind = op.kind
      if kind.endswith("_right"):
        kind = kind[:-len("_right")]  # order does not matter here
      if kind.endswith("_left"):
        kind = kind[:-len("_left")]

      import numpy
      from tensorflow.python.framework import tensor_util

      def _is_negative(x__):
        if isinstance(x__, numpy.ndarray):
          return (x__ < 0).any()
        if isinstance(x__, (int, float, numpy.number)):
          return x__ < 0
        assert isinstance(x__, tf.Tensor)
        x__ = tensor_util.constant_value(x__)
        if x__ is not None:
          return _is_negative(x__)
        return False

      def _bin_op(a, b):
        if a is None or b is None:
          return None
        with tf_util.same_control_flow_ctx([a, b]):
          if kind == "add":
            use_relu = _is_negative(a) or _is_negative(b)  # for dynamic tensors, assume all positive
            if use_relu:
              return tf.convert_to_tensor(tf_util.simplify_non_negative_seq_length(a + b))
            return a + b
          elif kind == "sub":
            return tf.convert_to_tensor(tf_util.simplify_non_negative_seq_length(a - b))
          elif kind == "mul":
            return a * b
          elif kind in ("floordiv", "truediv"):  # truediv assumes there is no remainder
            return a // b
          elif kind == "ceildiv":
            return -(-a // b)
          else:
            raise ValueError("unknown op kind %r" % op.kind)

      y_name = self.description + ":seq-length"
      y = None
      inputs = list(op.inputs)
      assert inputs
      while inputs:
        x = inputs.pop(0)
        if x.dimension is not None:
          if y is None:
            with tf.control_dependencies(None):  # this will reset the context
              y = Data(
                name=y_name, dim_tags=[], dtype="int32",
                placeholder=tf.constant(x.dimension))
            continue
          y.placeholder = _bin_op(y.placeholder, x.dimension)
          continue
        if self.batch:
          x = x.get_for_batch_ctx(self.batch, self.control_flow_ctx)
        x.complete_dyn_size()
        if not x.dyn_size_ext:
          return
        x = x.dyn_size_ext
        if y is None:
          y = x.copy(name=y_name)
          continue
        if x.dim_tags != y.dim_tags:
          common = Data.get_common_data([x, y], allow_broadcast_all_sources=True)
          x_ = x.copy_compatible_to(common) if x.dim_tags else x
          y_ = y.copy_compatible_to(common) if y.dim_tags else y
          y = common
        else:
          x_, y_ = x, y
        y.placeholder = _bin_op(y_.placeholder, x_.placeholder)
      assert y
      if self.dyn_size_ext:
        assert self.dyn_size_ext.dim_tags == y.dim_tags
      if y.batch:
        if self.batch:
          assert self.batch == y.batch
        else:
          self.batch = y.batch
      self.dyn_size_ext = y
      if y.placeholder is not None:
        self.set_tag_on_size_tensor(y.placeholder)

  def is_equal(self, other, ignore_feature_dim=False, allow_same_feature_dim=False, allow_same_spatial_dim=None,
               treat_feature_as_spatial=False, broadcast_matches=False, unknown_spatial_matches=False,
               undefined_matches=False, derived_matches=False):
    """
    Compares self to other for equality.

    Note that the default behavior is very restrictive.
    Use functions such as :func:`get_all_dimension_tags` or :func:`get_existing_tag_from_collection`
    to explicitly specify the behavior for the comparison.

    Also note that the definition is slightly ad-hoc for some cases,
    and might potentially change in the future.
      https://github.com/rwth-i6/returnn/issues/634

    :param Dim other:
    :param bool ignore_feature_dim:
    :param bool allow_same_feature_dim:
    :param bool|None allow_same_spatial_dim:
    :param bool treat_feature_as_spatial:
    :param bool broadcast_matches:
    :param bool unknown_spatial_matches:
    :param bool undefined_matches:
    :param bool derived_matches:
    :rtype: bool
    """
    if self is other:  # first some fast path check
      return True
    if self.special or other.special:
      return False  # only true if same instance, check above
    if self.generic or other.generic:
      # Note that this invalidates the transitive property of equivalence relations.
      if self.generic and self.dimension:
        if not other.generic or other.dimension:
          if self.dimension != other.dimension:
            return False
      if other.generic and other.dimension:
        if not self.generic or self.dimension:
          if self.dimension != other.dimension:
            return False
      return self.kind == other.kind
    if allow_same_spatial_dim is None:
      allow_same_spatial_dim = allow_same_feature_dim
    self_base = self.get_same_derived_base() if derived_matches else self.get_same_base()
    other_base = other.get_same_derived_base() if derived_matches else other.get_same_base()
    if self_base is other_base:
      return True
    if self_base.derived_from_op and other_base.derived_from_op:
      if self_base.derived_from_op == other_base.derived_from_op:
        return True
    self_kind = self.kind
    other_kind = other.kind
    if self_kind == other_kind == self.Types.Feature and ignore_feature_dim:
      return True
    if treat_feature_as_spatial:
      if self_kind == self.Types.Feature:
        self_kind = self.Types.Spatial
      if other_kind == self.Types.Feature:
        other_kind = self.Types.Spatial
    if self.dimension != other.dimension:
      if broadcast_matches and (self.dimension == 1 or other.dimension == 1):
        pass  # pass on
      else:
        return False
    if self_kind != other_kind:
      return False
    if self_kind == other_kind == self.Types.Batch:
      # Note: This might be incorrect in some cases,
      # e.g. for beam search when we have the beam hidden in the batch dim,
      # or when we used MergeDimsLayer on the batch axis, or so.
      # We might need to extend the logic here later.
      return True
    if self_kind == other_kind == self.Types.Feature:
      if allow_same_feature_dim:
        return True
    if self_kind == other_kind == self.Types.Spatial:
      if allow_same_spatial_dim:
        if self.dimension is not None:
          return True
        if broadcast_matches and (self.dimension == 1 or other.dimension == 1):
          return True
      if unknown_spatial_matches and ((self.dyn_size is None) or (other.dyn_size is None)):
        return True
      if undefined_matches and (self.undefined or other.undefined):
        return True
    # In principle, we would want to check for identity (self is other).
    # We currently use the description because the identity would not be the same
    # in case of template construction where a dim tag is once created for a template layer,
    # and then later again for the real layer.
    if self.auto_generated and other.auto_generated and self.description == other.description:
      return True
    return False

  def __eq__(self, other):
    """
    :param Dim other:
    :rtype: bool
    :return: :func:`is_equal` with default options
    """
    if not isinstance(other, Dim):
      return False
    return self.is_equal(other)

  def __ne__(self, other):
    """
    :param Dim other:
    :rtype: bool
    """
    return not (self == other)

  def __hash__(self):
    """
    :rtype: int
    :return: hash, matching to :func:`__eq__` (ignoring generic flag)
    """
    # This must match the behavior in __eq__, which is is_equal with default options.
    # I.e. different hash implies not equal (but same hash not necessarily equal).
    if self.generic:
      raise ValueError(
        "Hash for generic dim tag %s is not well defined. " % self +
        "The generic flag invalidates the transitive property of equivalence relations. "
        "Explicitly go through the set or dict of dim tags and check each for equality instead.")
    if self.special:
      return hash(id(self))
    if self.is_batch_dim():
      return hash(())
    base = self.get_same_base()
    if base is not self:
      return hash(base)
    if self.derived_from_op:
      return hash(self.derived_from_op)
    if self.auto_generated:
      return hash((base.kind, base.dimension, base.description))
    return hash(id(base))

  def __lt__(self, other):
    """
    Define some order. This is just such that `sorted` works, or some diff reporting, or so.
    It is on symbolic level, i.e. it does not consider the actual dimension value.
    The defined order somewhat arbitrary, so do not rely on the exact behavior,
    as this might change at some later point.
    Currently, it depends on the creation index.

    :param Dim other:
    :rtype: bool
    """
    if not isinstance(other, (Dim, _MarkedDim)):
      raise TypeError("cannot compare %r with %r" % (self, other))
    if self == other:
      return False
    return _dim_cmp_value(self) < _dim_cmp_value(other)

  def __gt__(self, other):
    """
    See :func:`__lt__`.

    :param Dim other:
    :rtype: bool
    """
    return other < self

  def __ge__(self, other):
    return not self < other

  def __le__(self, other):
    return not self > other

  def get_same_base(self):
    """
    :rtype: Dim
    """
    base = self
    while base.same_as:
      base = base.same_as
    return base

  def get_same_derived_base(self):
    """
    :rtype: Dim
    """
    base = self
    visited = {}
    while base.same_as or base.derived_from_tag:
      assert id(base) not in visited  # should not have cycles. normally this should never be triggered
      visited[id(base)] = base
      if base.same_as:
        base = base.same_as
        continue
      base = base.derived_from_tag
      assert base
    return base

  def get_derived_bases_list(self):
    """
    :rtype: list[Dim]
    """
    res = [self]
    base = self
    visited = {}
    while base.same_as or base.derived_from_tag:
      assert id(base) not in visited  # should not have cycles. normally this should never be triggered
      visited[id(base)] = base
      if base.same_as:
        base = base.same_as
        continue
      base = base.derived_from_tag
      assert base
      res.append(base)
    return res

  @property
  def undefined(self):
    """
    :return: whether the undefined flag is set, in self, bases, or any derived bases. also see :func:`is_dim_known`
    :rtype: bool
    """
    base = self
    visited = {}
    while base.same_as or base.derived_from_tag:
      assert id(base) not in visited  # should not have cycles. normally this should never be triggered
      visited[id(base)] = base
      if base._undefined:
        return True
      if base.same_as:
        base = base.same_as
        continue
      base = base.derived_from_tag
      assert base
    return base._undefined

  def declare_same_as(self, other):
    """
    :param Dim other:
    """
    assert self.can_be_used_as_dim() and other.can_be_used_as_dim()  # declare_same_as does not make sense otherwise
    self._maybe_update()
    self._validate_in_current_graph()
    if self == other:
      return
    other_same_base = other.get_same_base()
    if self is other_same_base or self.same_as is other_same_base:
      return
    self_same_as = self.get_same_base()
    if self_same_as is other_same_base:
      return
    if other_same_base.get_same_derived_base() is self_same_as:
      # We actually want it to be the other way around.
      other_same_base.declare_same_as(self_same_as)
      return
    if self.is_dim_known() and not other.is_dim_known():
      if self_same_as._creation_idx < other_same_base._creation_idx:
        # We want to keep self instead.
        # https://github.com/rwth-i6/returnn_common/issues/200
        other_same_base.declare_same_as(self_same_as)
        return
    other_derived_bases = set(other.get_derived_bases_list())
    self_derived_bases = set(self.get_derived_bases_list())
    assert other_derived_bases != self_derived_bases
    if self_derived_bases.issubset(other_derived_bases):
      # Avoid cycles on derived_from_tag. https://github.com/rwth-i6/returnn/issues/1054
      return other.declare_same_as(self)
    if self_same_as is not self:
      assert not self_same_as.same_as
      if self_same_as is other_same_base:
        return
      self_same_as.declare_same_as(other_same_base)
      if (self.dyn_size_ext is None or not self._validate_in_current_graph()) and self_same_as.dyn_size_ext:
        self.dyn_size_ext = self_same_as.get_dyn_size_ext_for_batch_ctx(self.batch, self.control_flow_ctx)
    other_same_base._merge_same_for_batch_ctx_dict(self)
    other._maybe_update()
    self.same_as = other_same_base
    self._same_as_tb = traceback.extract_stack()
    self._maybe_update()
    if self.dyn_size is not None and other_same_base.dyn_size is not None:
      if self.dyn_size is not other_same_base.dyn_size:
        if self.batch == other_same_base.batch and self.control_flow_ctx == other_same_base.control_flow_ctx:
          # Note: Instead of making this a warning, we could also enforce this at some point.
          #   The user should be able to fix `extern_data` in the config such that this is correct in the first place.
          #   Also, in addition to this warning, we might want to add some runtime check on the eq of the dyn sizes.
          print(
            "Warning: assuming dim tags are same with different size placeholders: %r vs %r" % (
              self.dyn_size, other_same_base.dyn_size))
    # If we have a defined source, and this is a dynamic spatial axis, and it was undefined before,
    # maybe we can overtake the size_placeholder now.
    if other_same_base.dyn_size is not None and self.src_data:
      assert isinstance(self.src_axis, int)
      # Maybe it changed in the meanwhile, so check.
      tag = self.src_data.get_dim_tag(self.src_axis)
      if tag.description == self.description and (not tag.dyn_size_ext or not tag._validate_in_current_graph()):
        tag.dyn_size_ext = self.get_dyn_size_ext_for_batch_ctx(tag.batch, tag.control_flow_ctx)
        tag._maybe_update()
    # If others dyn_size is None but we have a dyn_size, maybe update others dyn_size.
    if self.dyn_size is not None and other_same_base.dyn_size is not self.dyn_size:
      # Could be unset if it comes from the config, or from prev graph creation.
      # This is important such that self.can_compare() is sane.
      if other_same_base.dyn_size is None or not other_same_base._validate_in_current_graph():
        other_same_base.dyn_size_ext = self.get_dyn_size_ext_for_batch_ctx(
          other_same_base.batch, other_same_base.control_flow_ctx)
        other_same_base._maybe_update()
    if not self.dyn_size_ext or not self._validate_in_current_graph():
      self.dyn_size_ext = other_same_base.get_dyn_size_ext_for_batch_ctx(self.batch, self.control_flow_ctx)
      self._maybe_update()
    elif other_same_base.dyn_size_ext is None or not other_same_base._validate_in_current_graph():
      other_same_base.dyn_size_ext = self.get_dyn_size_ext_for_batch_ctx(
        other_same_base.batch, other_same_base.control_flow_ctx)
      other_same_base._maybe_update()
    if self.is_dim_known() and other.is_dim_known():
      assert self.dimension == other.dimension
    elif self.is_dim_known() and not other.is_dim_known():
      other.dimension = self.dimension
    elif not self.is_dim_known() and other.is_dim_known():
      self.dimension = other.dimension
    if self._vocab and not other_same_base._vocab:
      other_same_base._vocab = self._vocab
    elif other_same_base._vocab and not self._vocab:
      self._vocab = other_same_base._vocab
    if self.derived_from_op and not other_same_base.derived_from_op:
      other_same_base.derived_from_op = self.derived_from_op
    elif other_same_base.derived_from_op and not self.derived_from_op:
      self.derived_from_op = other_same_base.derived_from_op

  def _merge_same_for_batch_ctx_dict(self, other):
    """
    :param Dim other:
    """
    self._validate_in_current_graph()
    for _, dim in list(self._same_for_batch_ctx.items()):
      assert isinstance(dim, Dim)
      dim._validate_in_current_graph()
    for key, dim in other._same_for_batch_ctx.items():
      if not dim._validate_in_current_graph():
        continue
      self_dim = self._same_for_batch_ctx.get(key, None)
      if self_dim and (self_dim.dyn_size_ext or not dim.dyn_size_ext):
        continue  # keep ours
      if not dim.dyn_size_ext:
        continue  # undefined, do not overtake
      self._same_for_batch_ctx[key] = dim
    other._same_for_batch_ctx.clear()  # we only want to have it once

  @classmethod
  def get_existing_tag_from_collection(cls, other, tags, is_equal_opts=None):
    """
    :param Dim other:
    :param list[Dim]|tuple[Dim]|set[Dim] tags:
    :param dict[str]|None is_equal_opts: passed to Dim.is_equal
    :rtype: Dim|None
    """
    if is_equal_opts is None:
      is_equal_opts = {}
    # We do potential multiple rounds, such that we prefer "more equal" (using less is_equal_opts).
    rounds = [{}]
    if is_equal_opts:
      if "broadcast_matches" in is_equal_opts:
        rounds.append({k: v for (k, v) in is_equal_opts.items() if k != "broadcast_matches"})
      rounds.append(is_equal_opts)
    for _is_equal_opts in rounds:
      for _tag in tags:
        if _tag.is_equal(other, **_is_equal_opts):
          return _tag
    return None

  @classmethod
  def get_all_dimension_tags(cls, data_list, is_equal_opts=None, unique_separate_axes=True):
    """
    :param list[Data] data_list:
    :param dict[str]|None is_equal_opts: passed to Dim.is_equal
    :param bool unique_separate_axes: e.g. data_list=[Data with shape (B,5,5,10)] results in 4 dim tags, not 3.
    :return: list of dimension tags, dict for data -> list of dimension tags (for each axis)
    :rtype: (list[Dim], dict[Data, list[Dim]])
    """
    tags = []
    data_axes_dict = {}
    for data in data_list:
      data_axes_dict[data] = []
      existing_tag_collection_for_data = list(tags) if unique_separate_axes else tags
      for axis in range(data.batch_ndim):
        tag = data.get_dim_tag(axis)
        existing_tag = cls.get_existing_tag_from_collection(
          tag, tags=existing_tag_collection_for_data, is_equal_opts=is_equal_opts)
        if existing_tag:
          if unique_separate_axes:
            existing_tag_collection_for_data.remove(existing_tag)  # don't take it again for this data
          replace_existing = existing_tag.undefined and not tag.undefined and tag.dimension == existing_tag.dimension
          if tag != existing_tag and tag in existing_tag.get_derived_bases_list():
            replace_existing = True
          if replace_existing:  # Replace the existing by the new tag.
            tags[tags.index(existing_tag)] = tag
            existing_tag = tag
        else:  # no existing tag
          tags.append(tag)
        data_axes_dict[data].append(existing_tag or tag)
    return tags, data_axes_dict

  @classmethod
  def get_uniq_collection(cls, tags, is_equal_opts=None):
    """
    :param list[Dim]|tuple[Dim]|set[Dim] tags:
    :param dict[str]|None is_equal_opts: passed to Dim.is_equal
    :rtype: list[Dim]
    """
    res = []
    for tag in tags:
      ex = cls.get_existing_tag_from_collection(tag, res, is_equal_opts=is_equal_opts)
      if not ex:
        res.append(tag)
    return res

  def get_dim_value(self):
    """
    Infers the dim this axis should have if unbroadcasted.
    If `self.src_data` has a placeholder, will use the shape from there.
    Otherwise, uses `self.dimension` (if static) or `self.dyn_size` (if dynamic).
    :rtype: int|tf.Tensor
    """
    if self.dimension is not None:
      return self.dimension
    if self.is_batch_dim():
      if self.src_data:
        return self.src_data.get_batch_dim()
      from returnn.tf.layers.base import LayerBase
      return LayerBase.get_recent_layer().get_batch_dim()
    if self.src_data is not None and self.src_axis is not None and self.src_data.placeholder is not None:
      from returnn.tf.util.basic import get_shape_dim
      return get_shape_dim(self.src_data.placeholder, self.src_axis)
    self.complete_dyn_size()
    if self.dyn_size is not None:
      return tf.math.reduce_max(self.dyn_size)
    raise Exception('%s: need placeholder, self.dimension or self.dyn_size for dim value' % self)

  def axis_split_info(self):
    """
    :return: axis split info. see :func:`get_param_axes_split_info` and usage (e.g. pretraining)
    :rtype: list[int|None]
    """
    same_base = self.get_same_base()
    op = self.derived_from_op or same_base.derived_from_op
    if not op:
      return [self.dimension]
    if op.kind == "add":
      return sum([x.axis_split_info() for x in op.inputs], [])  # flatten
    if op.kind == "mul":
      res = [1]
      for x in op.inputs:
        res = sum([n * x.axis_split_info() if n is not None else None for n in res], [])  # flatten
      return res
    return [self.dimension]

  @property
  def vocab(self):
    """
    :rtype: returnn.datasets.util.vocabulary.Vocabulary|None
    """
    return self.get_same_base()._vocab

  @vocab.setter
  def vocab(self, vocab):
    """
    :param returnn.datasets.util.vocabulary.Vocabulary|None vocab:
    """
    self.get_same_base()._vocab = vocab

  # The kind of operations we have:
  # a + b: padding, concat
  # a - b: window with valid frames only
  # a * b: merge dims, upsample, transposed conv with striding
  # a / b (when a % b == 0): split dims, downsample, conv with striding
  # ceildiv(a, b): conv with striding
  # custom: repeat, remove, mask, loop with dyn end
  # Note that we differentiate between the order, i.e. a + b != b + a.
  # Note that we always have the assumption that a dimension is non-negative.
  # This assumption is necessary for some simplifications.
  # https://github.com/rwth-i6/returnn/pull/853

  def __add__(self, other):
    """
    :param Dim|int other:
    :return: self + other. note that this is not commutative, i.e. different from other + self.
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_add_sub_(other, kind="add", right=True)
    return term.as_dim()

  def __radd__(self, other):
    """
    :param Dim|int other:
    :return: other + self
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_add_sub_(other, kind="add", right=False)
    return term.as_dim()

  def __sub__(self, other):
    """
    :param Dim|int other:
    :rtype: Dim
    """
    return self.sub_right(other)

  def sub_right(self, other):
    """
    :param Dim|int other:
    :return: self - other
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_add_sub_(other, kind="sub", right=True)
    return term.as_dim()

  def sub_left(self, other):
    """
    :param Dim|int other:
    :return: (-other) + self
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_add_sub_(other, kind="sub", right=False)
    return term.as_dim()

  def __mul__(self, other):
    """
    :param Dim|int other:
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_mul_div_(other, kind="mul", right=True)
    return term.as_dim()

  def __rmul__(self, other):
    """
    :param Dim|int other:
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_mul_div_(other, kind="mul", right=False)
    return term.as_dim()

  def __floordiv__(self, other):
    """
    :param Dim|int other:
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_mul_div_(other, kind="floordiv", right=True)
    return term.as_dim()

  def __truediv__(self, other):
    """
    :param Dim|int other:
    :rtype: Dim
    """
    return self.div_right(other)

  def div_left(self, other):
    """
    :param Dim|int other:
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_mul_div_(other, kind="truediv", right=False)
    return term.as_dim()

  def div_right(self, other):
    """
    :param Dim|int other:
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_mul_div_(other, kind="truediv", right=True)
    return term.as_dim()

  def ceildiv_left(self, other):
    """
    :param Dim|int other:
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_mul_div_(other, kind="ceildiv", right=False)
    return term.as_dim()

  def ceildiv_right(self, other):
    """
    :param Dim|int other:
    :rtype: Dim
    """
    term = Dim._OpLinearTerm.from_dim(self)
    term.extend_mul_div_(other, kind="ceildiv", right=True)
    return term.as_dim()

  def __neg__(self):
    """
    :rtype: Dim
    """
    return -1 * self

  @classmethod
  def _make_constant_static_dim(cls, value, kind=None):
    """
    :param int value:
    :param Entity|None kind:
    :rtype: Dim
    """
    return Dim(
      dimension=value,
      kind=kind or Dim.Types.Unspecified,
      description="unnamed_%sdim_%i" % (kind.name + "_" if kind else "", value),
      derived_from_op=Dim.Op(kind="constant", inputs=[], attribs={"value": value}),
      auto_generated=True)

  def _is_constant_static_dim(self):
    return self.derived_from_op and self.derived_from_op.kind == "constant"

  class Op:
    """
    Op on :class:`Dim` which results in a derived :class:`Dim`.
    """
    def __init__(self, kind, inputs, attribs=None):
      """
      :param str kind: "add", "sub", "mul", "ceildiv"
      :param list[Dim] inputs:
      :param dict[str]|None attribs:
      """
      self.kind = kind
      self.inputs = inputs
      self.output = None  # type: typing.Optional[Dim]
      self.attribs = attribs

    def __repr__(self):
      attribs = (" %r" % self.attribs) if self.attribs else ""
      return "<Dim.Op %r %s%s>" % (self.kind, self.inputs, attribs)

    def _value(self):
      return self.kind, tuple(self.inputs), frozenset(self.attribs.items()) if self.attribs else None

    def __hash__(self):
      return hash(self._value())

    def __eq__(self, other):
      if isinstance(other, Dim.Op):
        return self._value() == other._value()
      return False

    def __ne__(self, other):
      return not self.__eq__(other)

  class _OpMultTerm:
    """
    represents sth like a * b * c
    """
    @classmethod
    def from_dim(cls, dim):
      """
      :param Dim dim:
      :rtype: Dim._OpMultTerm
      """
      dim = dim.get_same_base()
      if dim.dimension == 1 and dim._is_constant_static_dim():
        return cls.one()
      if dim.derived_from_op and dim.derived_from_op.kind == "mul":
        return cls(list(dim.derived_from_op.inputs))
      return cls([dim])

    @classmethod
    def from_dim_factors(cls, dims):
      """
      :param list[Dim] dims:
      :rtype: Dim._OpMultTerm
      """
      res = cls.one()
      for d in dims:
        res.extend_mul_div_(d, kind="mul", right=True)
      return res

    @classmethod
    def one(cls):
      """
      :rtype: Dim._OpMultTerm
      """
      return cls([])

    def __init__(self, terms):
      """
      :param list[Dim] terms:
      """
      self.terms = terms

    def __hash__(self):
      return hash(tuple(self.terms))

    def __eq__(self, other):
      """
      :param Dim|Dim._OpMultTerm other:
      """
      if isinstance(other, Dim._OpMultTerm):
        return self.terms == other.terms
      return False

    def __ne__(self, other):
      return not self.__eq__(other)

    def __repr__(self):
      return "Dim._OpMultTerm(%r)" % (self.terms,)

    @property
    def dimension(self):
      """
      :rtype: int|None
      """
      dim = 1
      for part in self.terms:
        if part.dimension is None:
          return None
        dim *= part.dimension
      return dim

    def base_term(self):
      """
      :rtype: Dim
      """
      assert self.terms
      return self.terms[-1]

    def is_one(self):
      """
      :rtype: bool
      """
      return not self.terms

    def is_constant_static_dim(self):
      """
      :rtype: bool
      """
      if not self.terms:
        return True
      return all(term._is_constant_static_dim() for term in self.terms)

    def copy(self):
      """
      :rtype: Dim._OpMultTerm
      """
      return Dim._OpMultTerm(list(self.terms))

    def negative(self):
      """
      :rtype: Dim._OpMultTerm
      """
      if self.terms and self.terms[0]._is_constant_static_dim() and self.terms[0].dimension == -1:
        return Dim._OpMultTerm(self.terms[1:])
      res = self.copy()
      res.extend_mul_div_(Dim._make_constant_static_dim(-1), kind="mul", right=False)
      return res

    def divisible(self, other, right):
      """
      :param Dim other:
      :param bool right:
      :return: whether we can divide other, without remainder
      :rtype: bool
      """
      if not self.terms:
        return False
      if other.derived_from_op and other.derived_from_op.kind == "mul":
        tmp = self.copy()
        for term in other.derived_from_op.inputs if right else reversed(other.derived_from_op.inputs):
          if not tmp.divisible(term, right=right):
            return False
          tmp.extend_mul_div_(term, kind="truediv", right=right)
        return True
      most_recent_term = self.terms[-1 if right else 0]
      if other == most_recent_term:
        return True
      if most_recent_term.dimension is not None and other.dimension is not None:
        if most_recent_term.dimension % other.dimension == 0:
          return True
      return False

    def can_simplify(self, other, kind, right):
      """
      :param Dim other:
      :param str kind:
      :param bool right:
      :return: whether we can simplify when applying this operation
      :rtype: bool
      """
      if other.derived_from_op and other.derived_from_op.kind == "mul":
        tmp = self.copy()
        for term in other.derived_from_op.inputs if right else reversed(other.derived_from_op.inputs):
          if not tmp.can_simplify(term, kind=kind, right=right):
            return False
          tmp.extend_mul_div_(term, kind=kind, right=right)
        return True
      idx = self._simplify_term_idx(other, kind=kind, right=right)
      return idx is not None

    def _simplify_term_idx(self, other, kind, right):
      """
      :param Dim other:
      :param str kind:
      :param bool right:
      :return: index of term to simplify
      :rtype: int|None
      """
      if not self.terms:
        return None
      for i, term in reversed(list(enumerate(self.terms))) if right else enumerate(self.terms):
        assert isinstance(term, Dim)
        if kind.endswith("div") and other == term:
          return i
        if kind == "mul" and term.derived_from_op:
          if term.derived_from_op.kind == "truediv_" + ("right" if right else "left"):
            if term.derived_from_op.inputs[-1] == other:
              return i
        if kind == "mul" and other.derived_from_op:
          if other.derived_from_op.kind == "truediv_" + ("right" if not right else "left"):
            if other.derived_from_op.inputs[-1] == term:
              return i
        if term._is_constant_static_dim() and other._is_constant_static_dim():
          if kind == "mul":
            return i
          if kind.endswith("div") and term.dimension % other.dimension == 0:
            return i
      return None

    def extend_mul_div_(self, other, kind, right):
      """
      :param Dim other:
      :param str kind:
      :param bool right:
      """
      assert kind in {"mul", "floordiv", "truediv", "ceildiv"}
      if other._is_constant_static_dim() and other.dimension == 1:
        return
      if not self.terms:
        if kind == "mul":
          self.terms.append(other)
        elif kind.endswith("div"):
          self.terms = [Dim._OpMultTerm.new_div_dim(self.as_dim(), other, kind=kind, right=right)]
        return
      if other.derived_from_op and other.derived_from_op.kind == "mul":
        for term in other.derived_from_op.inputs if right else reversed(other.derived_from_op.inputs):
          self.extend_mul_div_(term, kind=kind, right=right)
        return
      idx = self._simplify_term_idx(other, kind=kind, right=right)
      if idx is not None:
        term = self.terms[idx]
        assert isinstance(term, Dim)
        if kind.endswith("div") and other == term:
          self.terms.pop(idx)
          return
        if kind == "mul" and term.derived_from_op:
          if term.derived_from_op.kind == "truediv_" + ("right" if right else "left"):
            if term.derived_from_op.inputs[-1] == other:
              self.terms[idx] = term.derived_from_op.inputs[0]
              return
        if kind == "mul" and other.derived_from_op:
          if other.derived_from_op.kind == "truediv_" + ("right" if not right else "left"):
            if other.derived_from_op.inputs[-1] == term:
              self.terms[idx] = other.derived_from_op.inputs[0]
              return
        if term._is_constant_static_dim() and other._is_constant_static_dim():
          if kind == "mul":
            if term.dimension * other.dimension == 1:
              self.terms.pop(idx)
              return
            self.terms[idx] = Dim._make_constant_static_dim(
              term.dimension * other.dimension, kind=term.kind)
            return
          if kind.endswith("div") and term.dimension % other.dimension == 0:
            self.terms[idx] = Dim._make_constant_static_dim(
              term.dimension // other.dimension, kind=term.kind)
            return
          # Fallback with generic handling.
      if kind.endswith("div"):
        self.terms = [Dim._OpMultTerm.new_div_dim(self.as_dim(), other, kind=kind, right=right)]
        return
      if kind == "mul":
        if right:
          self.terms.append(other)
        else:
          self.terms.insert(0, other)
        return
      assert False

    @classmethod
    def new_div_dim(cls, numerator, denominator, kind, right):
      """
      :param Dim numerator:
      :param Dim denominator:
      :param str kind: "floordiv" or "ceildiv" or "truediv"
      :param bool right:
      :rtype: Dim
      """
      dim_value = None
      a = numerator.dimension
      b = denominator.dimension
      if a is not None and b is not None:
        if kind == "floordiv":
          dim_value = a // b
        elif kind == "ceildiv":
          dim_value = -(-a // b)
          if a % b == 0 and right:
            kind = "floordiv"  # for nicer description, and does not matter
        elif kind == "truediv":
          if a % b != 0:
            raise ValueError(
              "%s truediv %s only allowed if the result is an integer" % (numerator, denominator))
          dim_value = a // b
          if right:
            kind = "floordiv"  # for nicer description, and does not matter
        else:
          raise ValueError("invalid kind %r" % (kind,))
      if kind == "floordiv" and right:
        description = "%s//%s" % (Dim._get_description(numerator), Dim._get_description(denominator))
      else:
        description = "%s_%s(%s, %s)" % (
          kind, "right" if right else "left",
          Dim._get_description(numerator, brackets=False), Dim._get_description(denominator, brackets=False))
      op_kind = kind
      if a is not None and b is not None and a % b == 0:
        op_kind = "truediv"  # makes some other checks simpler
      op_kind += "_" + ("right" if right else "left")
      return Dim(
        description=description,
        kind=numerator.kind,
        dimension=dim_value,
        derived_from_op=Dim.Op(kind=op_kind, inputs=[numerator, denominator]),
        derived_from_tag=numerator)

    def as_dim(self):
      """
      :rtype: Dim
      """
      if self.is_one():
        return Dim._make_constant_static_dim(1)
      if len(self.terms) == 1:
        return self.terms[0]
      dim_kind = _get_merged_dim_kind(self.terms)
      return Dim(
        kind=dim_kind, description="*".join(map(Dim._get_description, self.terms)),
        dimension=self.dimension,
        derived_from_op=Dim.Op(kind="mul", inputs=list(self.terms)))

  @classmethod
  def _get_description(cls, dim, brackets=True):
    """
    :param Dim dim:
    :param bool brackets: add brackets when necessary
    :rtype: str
    """
    if dim.description and dim.description.startswith("unnamed_") and dim.dimension is not None:
      return str(dim.dimension)
    if dim.description:
      if brackets:
        import re
        if re.search("[+\\-/ ]", dim.description):
          return "(%s)" % dim.description
      return dim.description
    return "unnamed_%s_dim%s" % (dim.kind, dim.dimension if dim.dimension is not None else "?")

  class _OpLinearTerm:
    """
    represents sth like a * b + c
    """
    @classmethod
    def from_dim(cls, dim):
      """
      :param Dim dim:
      :rtype: Dim._OpLinearTerm
      """
      res = cls.zero()
      res.extend_add_sub_(dim, kind="add", right=True)
      return res

    @classmethod
    def zero(cls):
      """
      :rtype: Dim._OpLinearTerm
      """
      return Dim._OpLinearTerm([])

    def __init__(self, terms):
      """
      :param list[Dim._OpMultTerm] terms:
      """
      self.terms = terms

    def __hash__(self):
      return hash(tuple(self.terms))

    def __eq__(self, other):
      if isinstance(other, Dim._OpLinearTerm):
        return self.terms == other.terms
      return False

    def __ne__(self, other):
      return not self.__eq__(other)

    def as_dim(self):
      """
      :rtype: Dim
      """
      if self.is_zero():
        return Dim._make_constant_static_dim(0)
      if len(self.terms) == 1:
        return self.terms[0].as_dim()
      add_parts = []
      desc_parts = []
      dim = 0
      for term in self.terms:
        s = term.as_dim()
        add_parts.append(s)
        desc_parts.append(Dim._get_description(s))
        if dim is not None and s.dimension is not None:
          dim += s.dimension
        else:
          dim = None
      if len(add_parts) == 1:
        return add_parts[0]
      return Dim(
        kind=_get_merged_dim_kind(add_parts),
        description="+".join(desc_parts),
        dimension=dim,
        derived_from_op=Dim.Op(kind="add", inputs=add_parts),
        derived_from_tag=self.representative_tag())

    def __repr__(self):
      return "Dim._OpLinearTerm(%r)" % (self.terms,)

    def is_zero(self):
      """
      :rtype: bool
      """
      return not self.terms

    def extend_add_sub_(self, other, kind, right):
      """
      :param Dim|int other:
      :param str kind: "add" or "sub"
      :param bool right: or left. right means self + other, left means other + self
      """
      assert kind in {"add", "sub"}
      other = self._make_dim(other, kind=kind)
      if other._is_constant_static_dim() and other.dimension == 0:
        return
      if other.derived_from_op and other.derived_from_op.kind == "add":
        for other_ in other.derived_from_op.inputs if right else reversed(other.derived_from_op.inputs):
          self.extend_add_sub_(other_, kind=kind, right=right)
        return
      term = Dim._OpMultTerm.from_dim(other)
      neg_term = term.negative()
      if kind == "sub":
        term, neg_term = neg_term, term
      most_recent_term = self.terms[-1 if right else 0] if self.terms else None
      if most_recent_term:
        if most_recent_term == neg_term:
          self.terms.pop(-1 if right else 0)
          return
        if most_recent_term.is_constant_static_dim() and term.is_constant_static_dim():
          self.terms[-1 if right else 0] = Dim._OpMultTerm.from_dim(
            Dim._make_constant_static_dim(most_recent_term.dimension + term.dimension, kind=other.kind))
          return
        if most_recent_term.terms and term.terms and most_recent_term.terms[-1] == term.terms[-1]:
          # Merge terms
          a = Dim._OpMultTerm.from_dim_factors(most_recent_term.terms[:-1]).as_dim()
          b = Dim._OpMultTerm.from_dim_factors(term.terms[:-1]).as_dim()
          res = Dim._OpMultTerm.from_dim((a + b) if right else (b + a))
          res.extend_mul_div_(term.terms[-1], kind="mul", right=True)
          self.terms[-1 if right else 0] = res
          return
      if right:
        self.terms.append(term)
      else:
        self.terms.insert(0, term)

    def extend_mul_div_(self, other, kind, right):
      """
      :param Dim|int other:
      :param str kind: "mul" or "ceildiv"
      :param bool right: or left. right means self * other, left means other * self
      """
      assert kind in {"mul", "floordiv", "truediv", "ceildiv"}
      other = self._make_dim(other, kind=kind)
      if kind == "mul" and right:
        if not all(term.can_simplify(other, kind=kind, right=right) for term in self.terms):
          # Do it the other way around
          self.terms, other = Dim._OpLinearTerm.from_dim(other).terms, self.as_dim()
          right = False
      if other._is_constant_static_dim() and other.dimension == 1:
        return
      if kind.endswith("div"):
        if any(not term.divisible(other, right=right) for term in self.terms):
          self.terms = [Dim._OpMultTerm.from_dim(
            Dim._OpMultTerm.new_div_dim(self.as_dim(), other, kind=kind, right=right))]
          return
      for term in self.terms:
        term.extend_mul_div_(other, kind=kind, right=right)

    def _make_dim(self, other, kind):
      """
      :param Dim|int other:
      :param str kind:
      :rtype: Dim
      """
      if isinstance(other, int):
        base_tag = self.representative_tag()
        return Dim._make_constant_static_dim(other, kind=base_tag.kind if base_tag else None)
      elif isinstance(other, Dim):
        return other.get_same_base()
      else:
        raise TypeError("%s %s %s invalid for type %s" % (self, kind, other, type(other)))

    def representative_tag(self):
      """
      :rtype: Dim|None
      """
      # First find any dynamic.
      for term in self.terms:
        for term_ in term.terms:
          if term_.dimension is None:
            return term_
      # Now find non-unspecified.
      for term in self.terms:
        for term_ in term.terms:
          if term_.kind != Dim.Types.Unspecified:
            return term_
      # Now find any.
      for term in self.terms:
        for term_ in term.terms:
          return term_
      return None


# Earlier the class was called DimensionTag. Provide this alias for older code.
DimensionTag = Dim

# Global dim tag placeholders.
batch_dim = Dim(kind=Dim.Types.Batch, description="global batch")


# Provide some simple wrappers. https://github.com/rwth-i6/returnn/issues/782
# Use CamelCase function names (invalidates PEP8) to make it look like a class instance.

# noinspection PyPep8Naming
def FeatureDim(description, dimension, **kwargs):
  """
  :param str description:
  :param int|None dimension:
  :rtype: Dim
  """
  return Dim(kind=Dim.Types.Feature, description=description, dimension=dimension, **kwargs)


# noinspection PyPep8Naming
def SpatialDim(description, dimension=None, **kwargs):
  """
  :param str description:
  :param int|None dimension:
  :rtype: Dim
  """
  return Dim(kind=Dim.Types.Spatial, description=description, dimension=dimension, **kwargs)


# They match any feature or spatial dim respectively, as long as it is unique in a :class:`Data`.
any_feature_dim = FeatureDim("any-feature-dim", None, generic=True)
any_spatial_dim = SpatialDim("any-spatial-dim", None, generic=True)


# This indicates to perform a single step execution of some layer which can potentially have recurrent state.
single_step_dim = Dim(description="single-step", kind=Dim.Types.Spatial, special=True, dimension=1)


class _MarkedDim:
  def __init__(self, tag):
    """
    :param Dim tag:
    """
    self.tag = tag

  def __repr__(self):
    return "%s(%r)" % (self.__class__.__name__, self.tag)

  def _eq_tuple(self):
    return self.__class__, self.tag

  def __hash__(self):
    return hash(self._eq_tuple())

  def __eq__(self, other):
    if isinstance(other, _MarkedDim):
      return self._eq_tuple() == other._eq_tuple()
    return False

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

  def __lt__(self, other):
    """
    See :func:`Dim.__lt__`.
    """
    if not isinstance(other, (Dim, _MarkedDim)):
      raise TypeError("cannot compare %r with %r" % (self, other))
    if self == other:
      return False
    return _dim_cmp_value(self) < _dim_cmp_value(other)

  def __gt__(self, other):
    return other < self

  def __ge__(self, other):
    return not self < other

  def __le__(self, other):
    return not self > other


class _ImplicitDim(_MarkedDim):
  """
  Represents an implicit dim (dim tag) in :class:`Data`.
  https://github.com/rwth-i6/returnn/issues/706
  """


class ImplicitSparseDim(_ImplicitDim):
  """
  Represents an implicit dim via Data.sparse_dim.
  """


class ImplicitDynSizeDim(_ImplicitDim):
  """
  Represents an implicit dim via dynamic dim sizes.
  https://github.com/rwth-i6/returnn/issues/706
  (For example via :class:`CumConcatLayer`.)
  """


class OptionalDim(_MarkedDim):
  """
  Represents a dim which might exist or not.
  """


class VerifyOutShapeException(Exception):
  """
  Exception via :func:`Data.verify_out_shape`.
  """


class BatchInfo:
  """
  A batched tensor is a tensor with batch dimension,
  i.e. consisting of multiple samples/sequences
  which are supposed to be totally independent of each other.

  The batch dim can consists out of one or more flattened "virtual" dims,
  which :class:`BatchInfo` keeps track of.
  This class provides some additional information
  about the batch dimension.
  Only one instance per different type of batch-dim is expected
  (i.e. `batch_info1 is batch_info2` <==> same batch info).

  When we pass data from the dataset to the network
  (in all cases (training, inference ...) via :class:`Runner` in the TF engine),
  we get a batch dimension due to the minibatch construction.
  This is a global batch size
  (usually dynamic, because every minibatch/step can have a different amount of samples,
  although we can also support static sizes, which is needed e.g. for TPUs)
  represented by :class:`BatchInfo.GlobalBatchDim`.

  When we do beam search (see :class:`SearchBeam`),
  we have multiple hypotheses per batch item,
  and thus a different batch dimension.

  We can also pack arrays (multiple sequences)
  (also referred to as "flattened tensors", "non-padded tensors", "ragged tensors").
  See e.g. :class:`FlattenBatchLayer` or :func:`flatten_with_seq_len_mask`.
  Also see :class:`tf.RaggedTensor` which also represents
  packed tensors (but only batch-major, although this is just a reinterpretation).
  We do not directly use :class:`tf.RaggedTensor` in :class:`Data`
  to have robust and reliable code (which expects :class:`tf.Tensor`).
  However, we maybe can make use of some of the functions in :mod:`tf.ragged`.
  """

  class VirtualDimBase(object):
    """
    Represents one virtual dim, flattened into the batch dim.
    """
    def short_repr(self):
      """
      :rtype: str
      """
      raise NotImplementedError

    def __repr__(self):
      return "%s{%s}" % (self.__class__.__name__, self.short_repr())

  class FixedDim(VirtualDimBase):
    """
    Represents a dim with fixed size.
    """
    def __init__(self, size, dim_tag=None):
      """
      :param tf.Tensor|int size:
      :param Dim|None dim_tag:
      """
      self.size = size
      self.dim_tag = dim_tag

    def short_repr(self):
      """
      :rtype: str
      """
      if isinstance(self.size, int):
        return "F(%i)" % self.size
      return "F(?)"

  class GlobalBatchDim(FixedDim):
    """
    Represents the global batch dim by the network (minibatch construction from the dataset).
    """
    def short_repr(self):
      """
      :rtype: str
      """
      if isinstance(self.size, int) and self.size >= 0:
        return "B(%i)" % self.size
      return "B"

  class BeamDim(FixedDim):
    """
    Represents a search beam.
    """
    def __init__(self, beam):
      """
      :param SearchBeam beam:
      """
      super(BatchInfo.BeamDim, self).__init__(size=beam.beam_size)
      self.beam = beam

    def short_repr(self):
      """
      :rtype: str
      """
      return "Beam{%r}(%s)" % (self.beam.name, self.size)

  class PaddedDim(FixedDim):
    """
    Represents a dim with variable size, which is flattened with padding (not packed) into the batch.
    """
    def __init__(self, dim_tag):
      """
      :param Dim dim_tag:
      """
      super(BatchInfo.PaddedDim, self).__init__(size=dim_tag.get_dim_value())
      self.dim_tag = dim_tag

    def short_repr(self):
      """
      :rtype: str
      """
      return "Padded{%r}" % self.dim_tag.description

  class PackedDim(VirtualDimBase):
    """
    Represents a dim with variable sizes, which is packed (un-padded) into the batch.
    Variable w.r.t. other dims (must be per batch entry).
    """
    def __init__(self, dim_tag, key_axes):
      """
      :param Dim dim_tag:
      :param list[BatchInfo.VirtualDimBase] key_axes:
        most common case would be [GlobalBatchDim(...)],
        but [GlobalBatchDim(...),BeamDim(...)] is also common.
      """
      self.dim_tag = dim_tag
      self.key_axes = key_axes

    @property
    def sizes(self):
      """
      :return: shape [B_flat]
      :rtype: tf.Tensor
      """
      assert self.dim_tag.dyn_size is not None
      return self.dim_tag.dyn_size

    def short_repr(self):
      """
      :rtype: str
      """
      return "Packed{%r}" % (self.dim_tag.description,)

  def __init__(self, base, new_dim, new_dim_index=None):
    """
    :param BatchInfo|None base:
      If this is extended or based on another batch.
      Except of the batch dim of the dataset minibatch,
      we would always have a base.
    :param BatchInfo.VirtualDimBase|None new_dim:
    :param int|None new_dim_index:

    In most cases, this constructor would probably not be used directly by the user.
    """
    self.base = base
    virtual_dims = list(base.virtual_dims) if base else []
    if new_dim:
      if new_dim_index is None:
        assert not virtual_dims
        new_dim_index = 0
      if new_dim_index < 0:
        assert new_dim_index == -1
        virtual_dims.append(new_dim)
      else:
        virtual_dims.insert(new_dim_index, new_dim)
    self.virtual_dims = virtual_dims  # type: typing.List[BatchInfo.VirtualDimBase]
    self._dim = None  # type: typing.Optional[typing.Union[tf.Tensor,int]]
    # These self._global_... attributes are meant
    # to be accessed only via the global (root) object (via get_global_base).
    # They store global information.
    # We don't use class attributes because this should not be global per process but only per network.
    self._global_beam_dims_by_beam_name = {}  # type: typing.Dict[str,BatchInfo.BeamDim]
    self._global_padded_dims_by_dim_tag = {}  # type: typing.Dict[Dim,BatchInfo.PaddedDim]
    self._packed_dims_by_dim_tag = {}  # type: typing.Dict[Dim,BatchInfo.PackedDim]
    self.descendants = []  # type: typing.List[BatchInfo]
    self._descendants_by_beam_name = {}  # type: typing.Dict[str,BatchInfo]
    self._global_descendants_by_virtual_dims = {}  # type: typing.Dict[typing.Tuple[BatchInfo.VirtualDimBase,...],BatchInfo]  # noqa
    if base:
      base.descendants.append(self)
      if isinstance(new_dim, BatchInfo.BeamDim):
        beam = new_dim.beam
        assert beam.name not in base._descendants_by_beam_name
        base._descendants_by_beam_name[beam.name] = self
    global_base = self.get_global_base()
    assert tuple(self.virtual_dims) not in global_base._global_descendants_by_virtual_dims
    global_base._global_descendants_by_virtual_dims[tuple(self.virtual_dims)] = self

  # noinspection PyShadowingNames
  @classmethod
  def make_global_batch_info(cls, batch_dim):
    """
    :param tf.Tensor|int batch_dim:
    :return: global batch info w.r.t. the network / graph
    :rtype: BatchInfo
    """
    # This is not stored in a class attrib because this is only w.r.t. the network, not global in the process.
    return BatchInfo(base=None, new_dim=BatchInfo.GlobalBatchDim(size=batch_dim))

  _global_broadcast_batch = None  # type: typing.Optional[BatchInfo]

  @classmethod
  def make_global_broadcast_batch_info(cls):
    """
    :return: BatchInfo with no virtual dims, s.t. the dimension is 1 (== prod([])) (broadcastable)
    :rtype: BatchInfo
    """
    if cls._global_broadcast_batch:
      return cls._global_broadcast_batch
    cls._global_broadcast_batch = BatchInfo(base=None, new_dim=None)
    return cls._global_broadcast_batch

  @classmethod
  def get_common_batch_info(cls, batches):
    """
    :param list[BatchInfo|None] batches:
    :rtype: BatchInfo|None
    """
    # Fast paths.
    if not batches:
      return None
    if len(batches) == 1:
      return batches[0]
    # Make unique, and filter non-none.
    batches_ = []
    for batch in batches:
      if batch and batch not in batches_:
        batches_.append(batch)
    batches = batches_
    if not batches_:
      return None
    if len(batches) == 1:
      return batches[0]
    base = batches[0].get_global_base()

    # Collect all dims.
    all_virtual_dims = []
    for batch in batches:
      for dim in batch.virtual_dims:
        if dim not in all_virtual_dims:
          # We want to get a reasonable order.
          same_type_last_idx = None
          for i, dim_ in enumerate(all_virtual_dims):
            if type(dim_) == type(dim):
              same_type_last_idx = i
          if same_type_last_idx is not None:
            all_virtual_dims.insert(same_type_last_idx + 1, dim)
          else:
            all_virtual_dims.append(dim)

    # Check if some batch already has them all.
    for batch in batches:
      if set(batch.virtual_dims) == set(all_virtual_dims):  # allow different order here
        return batch

    # Ok, need to extend.
    global_batch_dims = [dim for dim in all_virtual_dims if isinstance(dim, BatchInfo.GlobalBatchDim)]
    assert len(global_batch_dims) == 1
    global_batch_dim = global_batch_dims[0]
    assert base.virtual_dims == [global_batch_dim]
    beams = [dim for dim in all_virtual_dims if isinstance(dim, BatchInfo.BeamDim)]
    if beams:
      base = base.copy_extend_with_beam(SearchBeam.get_combined_beam(*(b.beam for b in beams)))
    dim_idx = 0
    for dim in all_virtual_dims:
      if dim in global_batch_dims:
        dim_idx += 1 + len(beams)
        continue
      if dim in beams:
        continue
      base = base._copy_extend_dim(new_dim=dim, new_dim_idx=dim_idx)
      dim_idx += 1
    return base

  def __repr__(self):
    return "BatchInfo{%s}" % ", ".join([dim.short_repr() for dim in self.virtual_dims])

  def short_repr(self):
    """
    :rtype: str
    """
    # "x" is the Theano-style shortcut for a broadcast dim.
    return "&".join([dim.short_repr() for dim in self.virtual_dims] or ["Bx"])

  def __getstate__(self):
    raise Exception("Pickling of BatchInfo is not supported. (%s)" % self)

  @property
  def dim(self):
    """
    :rtype: tf.Tensor|int
    """
    if self._dim is not None:
      return self._dim
    if not self.virtual_dims:
      return 1
    if len(self.virtual_dims) == 1:
      dim = self.virtual_dims[0]
      assert isinstance(dim, BatchInfo.FixedDim)
      return dim.size
    from returnn.tf.util.basic import same_control_flow_ctx, optional_mul
    if all(isinstance(dim, BatchInfo.FixedDim) for dim in self.virtual_dims):
      dims = self.virtual_dims  # type: typing.List[BatchInfo.FixedDim]
      sizes = [dim.size for dim in dims]
      with same_control_flow_ctx(sizes):
        value = optional_mul(*sizes)  # type: typing.Union[tf.Tensor,int]
      self._dim = value
      return value
    if all(isinstance(dim, (BatchInfo.PackedDim, BatchInfo.GlobalBatchDim)) for dim in self.virtual_dims):
      dims = [dim for dim in self.virtual_dims if isinstance(dim, BatchInfo.PackedDim)]
      if len(dims) > 1:
        raise NotImplementedError("%s: currently only support one packed dim but have %r" % (self, dims))
      dim, = dims
      assert isinstance(dim, BatchInfo.PackedDim)
      with same_control_flow_ctx(dim.dim_tag.dyn_size_ext.placeholder):
        value = tf.reduce_sum(dim.dim_tag.dyn_size_ext.placeholder)
      self._dim = value
      return value
    raise NotImplementedError("%r.dim()" % self)

  @dim.setter
  def dim(self, value):
    """
    Can only set the global batch dim.

    :param tf.Tensor|int value:
    """
    assert len(self.virtual_dims) == 1
    dim = self.virtual_dims[0]
    assert isinstance(dim, BatchInfo.GlobalBatchDim)
    dim.size = value
    if dim.dim_tag:
      dim.dim_tag.dimension = value if (isinstance(value, int) and value > 0) else None
    self._dim = value

  @property
  def static_dim(self):
    """
    :rtype: int|None
    """
    # This should be safe. Do not call self.dim.
    if self._dim is not None:
      return self._dim if isinstance(self._dim, int) else None
    if not self.virtual_dims:
      return 1
    if len(self.virtual_dims) == 1:
      dim = self.virtual_dims[0]
      assert isinstance(dim, BatchInfo.FixedDim)
      return dim.size if isinstance(dim.size, int) else None
    from functools import reduce
    from operator import mul
    if all(isinstance(dim, BatchInfo.FixedDim) for dim in self.virtual_dims):
      dims = self.virtual_dims  # type: typing.List[BatchInfo.FixedDim]
      sizes = [dim.size for dim in dims]
      if all(isinstance(s, int) for s in sizes):
        return reduce(mul, sizes, 1)
      return None
    return None

  @property
  def beam(self):
    """
    :rtype: SearchBeam|None
    """
    beams = [dim for dim in self.virtual_dims if isinstance(dim, BatchInfo.BeamDim)]
    if beams:
      # Just return first. In case you need more custom logic, directly check the dims.
      return beams[0].beam
    return None

  def get_base_chain(self):
    """
    :rtype: list[BatchInfo]
    """
    bases = []
    base = self.base
    while base:
      bases.append(base)
      base = base.base
    return bases

  def get_global_base(self):
    """
    :rtype: BatchInfo
    """
    if not self.base:
      return self
    return self.get_base_chain()[-1]

  def get_global_batch_dim(self):
    """
    :rtype: BatchInfo.GlobalBatchDim
    """
    global_beam_dims = [dim for dim in self.virtual_dims if isinstance(dim, BatchInfo.GlobalBatchDim)]
    assert len(global_beam_dims) == 1
    return global_beam_dims[0]

  def is_global_batch(self):
    """
    :rtype: bool
    """
    global_beam_dims = [dim for dim in self.virtual_dims if isinstance(dim, BatchInfo.GlobalBatchDim)]
    return len(global_beam_dims) == 1 and len(self.virtual_dims) == 1

  def is_broadcast(self):
    """
    :rtype: bool
    """
    return len(self.virtual_dims) == 0

  def _make_beam_dim(self, beam):
    """
    :param SearchBeam beam:
    :rtype: BatchInfo.BeamDim
    """
    assert self.virtual_dims
    root = self.get_global_base()
    if beam.name in root._global_beam_dims_by_beam_name:
      return root._global_beam_dims_by_beam_name[beam.name]
    new_dim = BatchInfo.BeamDim(beam=beam)
    root._global_beam_dims_by_beam_name[beam.name] = new_dim
    return new_dim

  def _make_packed_dim(self, dim_tag):
    """
    :param Dim dim_tag:
    :rtype: BatchInfo.PackedDim
    """
    assert self.virtual_dims
    assert dim_tag.dyn_size is not None
    dim_tag_base = dim_tag.get_same_base()
    if dim_tag_base in self._packed_dims_by_dim_tag:
      return self._packed_dims_by_dim_tag[dim_tag_base]
    new_dim = BatchInfo.PackedDim(dim_tag=dim_tag, key_axes=self.virtual_dims)
    self._packed_dims_by_dim_tag[dim_tag_base] = new_dim
    return new_dim

  def _make_padded_dim(self, dim_tag):
    """
    :param Dim dim_tag:
    :rtype: BatchInfo.PaddedDim
    """
    assert self.virtual_dims
    root = self.get_global_base()
    assert dim_tag.dyn_size is not None
    dim_tag_base = dim_tag.get_for_batch_ctx(self, dim_tag.control_flow_ctx)
    if dim_tag_base in root._global_padded_dims_by_dim_tag:
      return root._global_padded_dims_by_dim_tag[dim_tag_base]
    new_dim = BatchInfo.PaddedDim(dim_tag=dim_tag_base)
    root._global_padded_dims_by_dim_tag[dim_tag_base] = new_dim
    return new_dim

  def _next_spatial_major_index(self):
    idx = None
    for i, dim in enumerate(self.virtual_dims):
      if isinstance(dim, BatchInfo.GlobalBatchDim):
        break
      if isinstance(dim, BatchInfo.BeamDim):
        break
      assert isinstance(dim, BatchInfo.FixedDim)
      idx = i + 1
    if idx is not None:
      return idx
    return 0

  def copy_extend_with_beam(self, beam):
    """
    :param SearchBeam beam:
    :rtype: BatchInfo
    """
    assert self.virtual_dims
    if self.beam == beam:
      return self
    if beam.name in self._descendants_by_beam_name:
      return self._descendants_by_beam_name[beam.name]
    return BatchInfo(
      base=self,
      new_dim=self._make_beam_dim(beam),
      new_dim_index=self.virtual_dims.index(self.get_global_batch_dim()) + 1)

  def copy_remove_beam(self):
    """
    :rtype: BatchInfo
    """
    if not self.beam:
      return self
    assert self.virtual_dims
    root = self.get_global_base()
    dims_wo_beam = [dim for dim in self.virtual_dims if not isinstance(dim, BatchInfo.BeamDim)]
    return root._global_descendants_by_virtual_dims[tuple(dims_wo_beam)]  # must exist

  def copy_remove_dim(self, remove_dim):
    """
    :param VirtualDimBase remove_dim:
    :rtype: BatchInfo
    """
    assert self.virtual_dims
    root = self.get_global_base()
    dims_wo_dim = [dim for dim in self.virtual_dims if dim != remove_dim]
    return root._global_descendants_by_virtual_dims[tuple(dims_wo_dim)]  # must exist

  def copy_set_beam(self, beam):
    """
    :param SearchBeam|None beam:
    :rtype: BatchInfo
    """
    batch = self.copy_remove_beam()
    if beam:
      batch = batch.copy_extend_with_beam(beam)
    return batch

  def copy_extend_with_packed_dim_tag(self, dim_tag, batch_major):
    """
    :param Dim dim_tag:
    :param bool batch_major: if True, add new dim in front. otherwise, add new dim at the end
    :rtype: BatchInfo
    """
    new_dim = self._make_packed_dim(dim_tag)
    new_dim_idx = -1 if batch_major else self._next_spatial_major_index()
    return self._copy_extend_dim(new_dim=new_dim, new_dim_idx=new_dim_idx)

  def copy_extend_with_padded_dim_tag(self, dim_tag, batch_major=None, new_dim_idx=None):
    """
    :param Dim dim_tag:
    :param bool|None batch_major: if True, add new dim in front. otherwise, add new dim at the end
    :param int|None new_dim_idx:
    :rtype: BatchInfo
    """
    new_dim = self._make_padded_dim(dim_tag)
    if new_dim_idx is None:
      assert batch_major is not None
      new_dim_idx = -1 if batch_major else self._next_spatial_major_index()
    else:
      assert batch_major is None
    return self._copy_extend_dim(new_dim=new_dim, new_dim_idx=new_dim_idx)

  def copy_extend_with_padded_or_fixed_dim_tag(self, dim_tag, batch_major=None, new_dim_idx=None):
    """
    :param Dim dim_tag:
    :param bool|None batch_major: if True, add new dim in front. otherwise, add new dim at the end
    :param int|None new_dim_idx:
    :rtype: BatchInfo
    """
    if dim_tag.dyn_size is not None:
      new_dim = self._make_padded_dim(dim_tag)
    else:
      new_dim = BatchInfo.FixedDim(size=dim_tag.get_dim_value(), dim_tag=dim_tag)
    if new_dim_idx is None:
      assert batch_major is not None
      new_dim_idx = -1 if batch_major else self._next_spatial_major_index()
    else:
      assert batch_major is None
    return self._copy_extend_dim(new_dim=new_dim, new_dim_idx=new_dim_idx)

  def _copy_extend_dim(self, new_dim, new_dim_idx):
    """
    :param BatchInfo.VirtualDimBase new_dim:
    :param int new_dim_idx:
    :rtype: BatchInfo
    """
    assert self.virtual_dims
    root = self.get_global_base()
    virtual_dims = list(self.virtual_dims)
    if new_dim_idx < 0:
      assert new_dim_idx == -1
      virtual_dims.append(new_dim)
    else:
      virtual_dims.insert(new_dim_idx, new_dim)
    if tuple(virtual_dims) in root._global_descendants_by_virtual_dims:
      return root._global_descendants_by_virtual_dims[tuple(virtual_dims)]
    return BatchInfo(base=self, new_dim=new_dim, new_dim_index=new_dim_idx)


class SearchBeam:
  """
  Represents info about the beam from some beam search (e.g. via :func:`beam_search`),
  e.g. such as the beam size, but also the dependencies.
  This is somewhat parallel to :class:`SearchChoices`, but simpler,
  and independent from the layers/network (:class:`returnn.tf.layers.base.LayerBase`).
  """

  def __init__(self, beam_size, dependency=NotSpecified, name=None, _next_frame=None):
    """
    :param int beam_size:
    :param SearchBeam|NotSpecified|None dependency:
    :param str|None name:
    :param SearchBeam|None _next_frame:
    """
    if isinstance(dependency, SearchBeam):
      assert name and dependency.name and name != dependency.name
    if name and os.path.basename(name).startswith("prev:"):
      assert _next_frame
    self.beam_size = beam_size
    self.dependency = dependency
    self.name = name
    self._next_frame = _next_frame

  def copy_as_prev_frame(self):
    """
    :rtype: SearchBeam
    """
    if self._next_frame:  # already prev frame -> return self. see logic in RecLayer maybe_transform
      return self
    assert self.name
    name = "%s/prev:%s" % (os.path.dirname(self.name), os.path.basename(self.name))
    return SearchBeam(beam_size=self.beam_size, name=name, _next_frame=self)

  def __repr__(self):
    keys = ["name", "beam_size"]
    if self.dependency is not NotSpecified:
      keys.append("dependency")
    return "%s(%s)" % (
      self.__class__.__name__, ", ".join(["%s=%r" % (key, getattr(self, key)) for key in keys]))

  def __eq__(self, other):
    """
    :param SearchBeam|object|None other:
    :rtype: bool
    """
    if self is other:
      return True
    if self is None or other is None:
      return False
    if not isinstance(self, SearchBeam) or not isinstance(other, SearchBeam):
      return False
    if self.name is None or other.name is None:
      return False  # cannot identify
    return self.name == other.name

  def __ne__(self, other):
    """
    :param SearchBeam|object|None other:
    :rtype: bool
    """
    return not (self == other)

  def __hash__(self):
    return hash(self.name)

  def _get_dependency_list(self):
    """
    :return: list as far as it is defined
    :rtype: list[SearchBeam]
    """
    ls = [self]
    while isinstance(ls[-1].dependency, SearchBeam):
      ls.append(ls[-1].dependency)
    return ls

  @classmethod
  def get_combined_beam(cls, beam1, beam2=None, *beams):
    """
    Combines beams.
    This will throw an exception if they cannot be combined.
    Note that in beam search (see :class:`SearchChoices`),
    the logic to combine beams from different search choices
    happens in a generic way for all layers automatically
    via :func:`TFNetwork._create_layer_layer_desc`,
    so normally we already have the same beam.
    Unless we are at template construction.

    :param SearchBeam|None beam1:
    :param SearchBeam|None beam2:
    :param SearchBeam|None beams:
    :rtype: SearchBeam|None
    """
    if beams:
      beam12 = cls.get_combined_beam(beam1, beam2)
      return cls.get_combined_beam(beam12, beams[0], *beams[1:])
    if beam2 is None:
      return beam1
    if beam1 is None:
      return beam2
    if beam1 == beam2:
      if beam2.dependency is NotSpecified:
        return beam1
      if beam1.dependency is NotSpecified:
        return beam2
      return beam1
    assert beam1.name and beam2.name
    if beam2._next_frame and not beam1._next_frame:
      return beam1
    if beam1._next_frame and not beam2._next_frame:
      return beam2
    b1 = beam1
    b2 = beam2
    used_next_frame = False
    if b1._next_frame and b2._next_frame:
      b1 = b1._next_frame
      b2 = b2._next_frame
      used_next_frame = True
    l1 = b1._get_dependency_list()
    l2 = b2._get_dependency_list()
    if b2 in l1:
      return beam1
    if b1 in l2:
      return beam2
    if used_next_frame:
      # Example: beam1: prev:out, beam2: prev:t, t->prev:out (l2).
      if beam1 in l2:  # -> beam1 dep on beam2
        return beam1
      if beam2 in l1:
        return beam2
    raise Exception(
      "\n".join([
        "Cannot combine beams:",
        "  1: %s (deps: %s, next %s, next deps %s)" % (
          beam1, beam1._get_dependency_list(),
          beam1._next_frame, beam1._next_frame._get_dependency_list() if beam1._next_frame else None),
        "  2: %s (deps: %s, next %s, next deps %s)" % (
          beam2, beam2._get_dependency_list(), beam2._next_frame,
          beam2._next_frame._get_dependency_list() if beam2._next_frame else None)]))


class Data(object):
  """
  This class is to describe a tensor,
  i.e. its shape and properties like
  whether we should consider it sparse data (i.e. it represents indices).
  Each dimension is described by :class:`Dim`.

  This is used in :class:`TFNetwork` to describe the dataset external data (:class:`ExternData`)
  as well as in every layer's output and in many other parts of the code.

  See :ref:`data`.
  """

  size_dtype = "int32"

  def __init__(self, name,
               shape=None, dtype=None,
               placeholder=None,
               sparse=None,
               sparse_dim=NotSpecified,
               dim=NotSpecified,
               size_placeholder=None,
               batch_dim_axis=NotSpecified,
               time_dim_axis=NotSpecified,
               feature_dim_axis=NotSpecified,
               available_for_inference=True,
               auto_create_placeholders=False,
               vocab=None,
               dim_tags=None,
               same_dim_tags_as=None,
               batch=None,
               beam=None,
               control_flow_ctx=None):
    """
    :param str name:
    :param tuple[int|None]|list[int|None] shape: including time-dim (can be None). excluding batch-dim.
      e.g. (time,feat)=(None,128)
    :param str dtype: e.g. "float32" or "int64"
    :param tf.Tensor|None placeholder: with added batch-dim
    :param bool|None sparse: whether to treat the value as an index. do not confuse with tf.SparseTensor
    :param Dim|int|None|NotSpecified sparse_dim:
    :param int|None|NotSpecified dim: feature dimension, shape[-1] if not sparse, otherwise like num_classes
    :param int|None|NotSpecified batch_dim_axis: where we add the batch-dim.
      e.g. shape=(time,...), 0 -> (batch,time,...), 1 -> (time,batch,...).
      Default is 0.
      This is normally always set, and a lot of code expects this. However, you can set it to None
      if this Data does not have a batch-dim.
    :param int|None|NotSpecified time_dim_axis: where we have the time dim axis, after we added the batch-dim.
      this is often 1. however, can be None if there is no time-dim.
    :param int|None|NotSpecified feature_dim_axis: feature dim axis. by default it's the last one
    :param dict[int,tf.Tensor]|None size_placeholder: for every None in shape, this will describe the size.
      The size is always a tensor of shape (batch,), i.e. the size can be different for each sequence in a batch.
    :param bool available_for_inference: e.g. the extern data "classes" is usually not available for inference
    :param bool auto_create_placeholders: This will create a tf.placeholder.
    :param str|dict[str]|returnn.datasets.util.vocabulary.Vocabulary|None vocab:
    :param tuple[Dim]|list[Dim]|dict[int,Dim]|None dim_tags:
      If tuple/list, this specifies the whole (batch) shape.
      If dict, explicitly specified dimension tags per axis (axis counted with batch-dim)
    :param dict[int|str,Dim]|None same_dim_tags_as: will mark our dimension tags to be the same
    :param BatchInfo|None batch:
    :param SearchBeam|None beam: the batch-dim could be extended by a beam-size,
      such that it represents the merged dims [batch, beam_size].
    :param ControlFlowContext|None control_flow_ctx:
    """
    assert isinstance(name, str)
    assert dtype is None or isinstance(dtype, str)
    self.name = name
    if sparse is None:
      sparse = sparse_dim not in (None, NotSpecified)
    if sparse_dim is NotSpecified:
      if sparse:
        assert dim is not NotSpecified, "need dim (num classes) if sparse"
        assert dim is None or isinstance(dim, int)
        sparse_dim = Dim(kind=Dim.Types.Feature, dimension=dim, description="%s:sparse-dim" % name, auto_generated=True)
      else:
        sparse_dim = None
    if isinstance(sparse_dim, int):
      sparse_dim = Dim(
        kind=Dim.Types.Feature, dimension=sparse_dim, description="%s:sparse-dim" % name, auto_generated=True)
    if sparse_dim is not None:
      assert isinstance(sparse_dim, Dim)
      assert sparse_dim.can_be_used_as_dim()
      assert sparse
      if dim is not NotSpecified:
        assert sparse_dim.dimension == dim
    else:
      assert not sparse
    self.sparse_dim = sparse_dim
    if dtype is None:
      if sparse:
        dtype = "int32"
      else:
        dtype = "float32"
    self.dtype = dtype  # type: str
    if beam and batch:
      assert batch.beam == beam
    self._batch = batch
    del batch
    self._beam = beam
    del beam
    self.control_flow_ctx = control_flow_ctx
    self.available_for_inference = available_for_inference
    if isinstance(dim_tags, (tuple, list)):
      assert all(tag.can_be_used_as_dim() for tag in dim_tags)
      # We do a couple of sanity checks, and maybe set special axes attribs.
      shape_ = tuple(tag.dimension for tag in dim_tags if not tag.is_batch_dim())
      if shape is not None:
        assert tuple(shape) == shape_
      del shape
      batch_dim_axis_ = _batch_dim_axis_from_dim_tags_tuple(dim_tags)
      if batch_dim_axis is not NotSpecified:
        assert batch_dim_axis == batch_dim_axis_
      del batch_dim_axis
      if time_dim_axis is NotSpecified:
        time_dim_axis = _default_time_dim_axis_dim_tags(dim_tags)
      dim_tags = tuple(dim_tags)
      if auto_create_placeholders:
        _auto_create_size_placeholders_on_dim_tags(name=name, dim_tags=dim_tags)
      if batch_dim_axis_ is not None:
        if dim_tags[batch_dim_axis_].batch and not self._batch:
          self._batch = dim_tags[batch_dim_axis_].batch
      del shape_
      del batch_dim_axis_
    else:
      if batch_dim_axis is NotSpecified:
        batch_dim_axis = 0
      if shape is None:
        if time_dim_axis is NotSpecified:
          time_dim_axis = _default_time_dim_axis_no_shape(
            batch_dim_axis=batch_dim_axis, feature_dim_axis=feature_dim_axis)
        shape, time_dim_axis = _infer_default_shape_and_time(
          batch_dim_axis=batch_dim_axis, feature_dim_axis=feature_dim_axis, time_dim_axis=time_dim_axis,
          sparse=sparse, dim=dim)
      else:
        if time_dim_axis is NotSpecified:
          time_dim_axis = _default_time_dim_axis(batch_dim_axis=batch_dim_axis, shape=shape)
      dim_tags = _infer_dim_tags_tuple_from_shape(
        shape, batch_dim_axis=batch_dim_axis, time_dim_axis=time_dim_axis, feature_dim_axis=feature_dim_axis,
        size_placeholder=size_placeholder, name=name,
        auto_create_placeholders=auto_create_placeholders,
        dim_tags=dim_tags, sparse=sparse)
      del batch_dim_axis
      del shape
    self._dim_tags = dim_tags  # type: typing.Tuple[Dim]
    if feature_dim_axis is not NotSpecified:
      if isinstance(feature_dim_axis, int):
        assert not self.sparse, "cannot have feature_dim_axis when sparse"
        if feature_dim_axis < 0:
          feature_dim_axis += self.batch_ndim
        assert 0 <= feature_dim_axis < self.batch_ndim
    self._feature_dim_axis = feature_dim_axis
    if time_dim_axis is not None:
      assert 0 <= time_dim_axis < self.batch_ndim
    self.time_dim_axis = time_dim_axis  # type: typing.Optional[int]  # counted with batch-dim
    if dim is not NotSpecified:
      if sparse:
        assert self.sparse_dim.dimension == dim
      else:
        if self.feature_dim_axis is None:
          assert dim is None
        else:
          assert self.batch_shape[self.feature_dim_axis] == dim
    if placeholder is None and auto_create_placeholders:
      with tf.name_scope("extern_data/placeholders/%s/" % name):
        placeholder = tf_compat.v1.placeholder(**self.get_placeholder_kwargs(with_batch=True))
    self._placeholder = placeholder  # type: tf.Tensor  # this will hold the data value itself
    if vocab is not None:
      from returnn.datasets.util.vocabulary import Vocabulary
      if isinstance(vocab, str):
        vocab = Vocabulary(vocab)
      elif isinstance(vocab, dict):
        vocab = Vocabulary.create_vocab(**vocab)
      assert isinstance(vocab, Vocabulary)
      assert self.sparse, "%s should represent indices of %s" % (self, vocab)
      assert self.dim == vocab.num_labels, "%s dims do not match with vocab %s" % (self, vocab)
      self.sparse_dim.vocab = vocab
    # The size_placeholder is for each variable length dimension in shape, i.e. excluding the batch-dim.
    if size_placeholder:
      self.size_placeholder = size_placeholder  # type: typing.Dict[int,tf.Tensor]  # axis w.o. batch -> size (batch,)
    if same_dim_tags_as:
      for _axis, _dim_tag in sorted(same_dim_tags_as.items()):
        _axis = self.get_axis_from_description(_axis)
        assert isinstance(_dim_tag, Dim)
        base_tag = self._dim_tags[_axis]
        if base_tag != _dim_tag:
          base_tag.declare_same_as(_dim_tag)
          if _dim_tag.dyn_size is not None:
            self.set_dynamic_size(_axis, _dim_tag.dyn_size)
    self._adapt_batch_consistent_dim_tags()
    self.sanity_check(assume_complete=False)

  @classmethod
  def from_tensor(cls, x):
    """
    :param tf.Tensor x:
    :rtype: Data
    """
    assert x.get_shape().is_fully_defined()
    x_shape = x.get_shape().as_list()
    return Data(name=str(x.op.name), shape=x_shape, batch_dim_axis=None, dtype=x.dtype.name, placeholder=x)

  @classmethod
  def template_from_constant(cls, x, name, dtype=None, shape=None, with_batch_dim=False, sparse_dim=None):
    """
    :param int|float|bool|numpy.ndarray x: not actually assigned to the returned Data, just for the shape and dtype
    :param str name:
    :param str|None dtype:
    :param list[Dim|int]|tuple[Dim|int]|None shape: for verification, and defining dim tags.
      might also already include the batch-dim. (Then with_batch_dim is ignored.)
    :param bool with_batch_dim:
    :param Dim|None sparse_dim:
    :return: data template
    :rtype: Data
    """
    import numpy
    if dtype is None:
      if isinstance(x, bool):
        dtype = "bool"
      elif isinstance(x, int):
        dtype = "int32"
      elif isinstance(x, float):
        dtype = "float32"
      elif isinstance(x, numpy.ndarray):
        dtype = str(x.dtype)
      else:
        raise TypeError("%r: cannot handle value %r of type %r" % (name, x, type(x)))
    shape_ = x.shape if isinstance(x, numpy.ndarray) else ()
    if shape is not None:
      if len(shape) > len(shape_) == 0:
        pass  # Scalar given, complex shape wanted. Allow this.
      else:
        assert len(shape) == len(shape_), "%r: shape does not match in ndim, %r vs %r" % (name, shape, shape_)
    else:
      shape = shape_
    dim_tags = []
    for i, d in enumerate(shape):
      d_ = shape_[i] if len(shape_) > 0 else None
      if isinstance(d, Dim):
        if len(shape_) > 0:
          assert d.dimension == d_
      elif isinstance(d, int):
        if len(shape_) > 0:
          assert d == d_
        d = Dim(
          kind=Dim.Types.Spatial if i < len(shape) - 1 else Dim.Types.Feature,
          description="%s:static:%i" % (name, i), auto_generated=True,
          dimension=d)
      else:
        raise TypeError("%r shape[%i] invalid type %r in shape %r" % (name, i, type(d), shape))
      dim_tags.append(d)
    if with_batch_dim and batch_dim not in dim_tags:
      dim_tags.insert(0, batch_dim)
    return Data(name=name, dim_tags=dim_tags, dtype=dtype, sparse_dim=sparse_dim)

  def sanity_check(self, ignore_placeholder=False, assume_complete=True):
    """
    Performs some sanity checks on self, and raises exceptions if something is not sane.

    :param bool ignore_placeholder:
    :param bool assume_complete:
    """
    for axis_name, axis in self.get_special_axes_dict().items():
      assert axis is None or 0 <= axis < self.batch_ndim, "%s: axis %s (%i) invalid" % (self, axis_name, axis)
    if self.batch_dim_axis is not None:
      for axis_name, axis in self.get_special_axes_dict().items():
        assert axis != self.batch_dim_axis, "%s: axis %s (%i) must be different from batch_dim_axis (%i)" % (
          self, axis_name, axis, self.batch_dim_axis)
    if self.sparse:
      assert self.feature_dim_axis is None, "%s: If sparse, there cannot be a feature dim axis." % self
    else:
      if self.feature_dim_axis is None:  # e.g. scalars, or [B]
        assert self.dim is None, "%s: not sparse but no feature-dim-axis, so dim should be None" % self
    if self.feature_dim_axis is not None:
      assert self.dim == self.batch_shape[self.feature_dim_axis], (
        "%s: inconsistent dim. feature axis or unspecified: %r." % (self, self.feature_dim_axis_or_unspecified))
    for axis, tag in enumerate(self.dim_tags):
      assert self.batch_shape[axis] == tag.dimension
      if tag.is_batch_dim():
        assert axis == self.batch_dim_axis, "%s: invalid %s" % (self, tag)
        continue  # further checks will assume not batch
      assert axis != self.batch_dim_axis, "%s: invalid %s" % (self, tag)
      # Note: tag.kind (feature or spatial) is independent from self.feature_dim_axis.
      if tag.batch and self.batch:
        assert tag.batch == self.batch or self.batch.is_broadcast()
      if tag.dyn_size_ext:
        assert tag.dyn_size_ext.dtype in {"int32", "int64"}
        if tag.dyn_size_ext.have_batch_axis():
          assert tag.batch == tag.dyn_size_ext.batch
        tag.dyn_size_ext.sanity_check()
    if not ignore_placeholder and self.placeholder is not None:
      # Note: We could just call self.placeholder.set_shape.
      # However, we are more explicit. We assume that the placeholder has already a known shape, and error otherwise.
      assert self.placeholder.shape.ndims == self.batch_ndim
      for i in range(self.batch_ndim):
        if self.batch_shape[i] is None:
          continue  # we allow anything in the placeholder
        if self.placeholder.shape[i].value != self.batch_shape[i]:
          print("Mismatching shape: Tensor %r vs Data %r" % (self.placeholder, self))
          from .basic import print_graph_output
          print_graph_output(self.placeholder, max_depth=3)
        assert self.placeholder.shape[i].value == self.batch_shape[i]
      self.placeholder.set_shape(self.batch_shape)
      assert self.placeholder.dtype.base_dtype.name == self.dtype
      # Currently only if placeholder is set.
      # We can later always do the check even without placeholder.
      if assume_complete:
        for tag in self.dim_tags:
          if tag.dimension is None:
            if tag.is_batch_dim():
              continue
            if not tag.dyn_size_ext or tag.dyn_size_ext.placeholder is None:
              tag.complete_dyn_size()
            assert tag.dyn_size is not None

  def get_runtime_sanity_check_op(self):
    """
    :return: op which does a couple of runtime sanity checks on the placeholder
    :rtype: tf.Operation
    """
    assert self.placeholder is not None
    checks = []
    with tf.name_scope("runtime_sanity_check"):
      shape = tf.shape(self.placeholder)
      # noinspection PyShadowingNames
      batch_dim = shape[self.batch_dim_axis] if self.have_batch_axis() else 1
      rank = tf.rank(self.placeholder)
      data = ["Data.get_runtime_sanity_check_op:", str(self), "shape", shape]
      for i, tag in enumerate(self.dim_tags):
        if tag.dyn_size is not None:
          data += [
            "dyn_size[%i] (%s)" % (i, tag), tag.dyn_size, ".shape", tf.shape(tag.dyn_size)]
      checks += [tf.Assert(tf.equal(rank, self.batch_ndim), data + ["-> invalid rank"])]
      if self.have_batch_axis():
        batch_dim_via_info = self.get_batch_dim()
        checks += [
          tf.Assert(tf.equal(batch_dim, batch_dim_via_info), data + ["-> invalid batch dim info", batch_dim_via_info])]
      for i in range(self.batch_ndim):
        if self.batch_shape[i] is not None:
          checks += [tf.Assert(tf.equal(shape[i], self.batch_shape[i]), data + ["-> invalid shape[%i]" % i])]
        dyn_size_ext = self.dim_tags[i].dyn_size_ext
        if dyn_size_ext and dyn_size_ext.placeholder is not None:
          dyn_size = dyn_size_ext.placeholder
          if dyn_size_ext.have_batch_axis() and self.have_batch_axis():
            checks += [tf.Assert(
              tf.equal(tf.shape(dyn_size)[dyn_size_ext.batch_dim_axis], batch_dim),
              data + ["-> invalid axis %i tag dyn size batch dim" % i])]
          checks += [tf.Assert(
            # Note: in almost all cases, we have equality here.
            # However, not strictly in all cases, e.g. DecideLayer, maybe some others...
            # But that should not be more than 1 less.
            tf.logical_or(
              tf.logical_and(
                tf.less_equal(tf.reduce_max(dyn_size), shape[i]),
                tf.greater_equal(tf.reduce_max(dyn_size), shape[i] - 1)),
              # In other rare cases, this might be a broadcast dim
              # (e.g. as initial values of att weights for a rec loop).
              tf.equal(1, shape[i])),
            data + ["-> invalid shape[%i] or max(dyn_size[%i])" % (i, i)])]
          checks += [dyn_size_ext.get_runtime_sanity_check_op()]
    return tf.group(*checks)

  def verify_out_shape(self, out_shape):
    """
    Verifies that ``out_shape`` matches our shape, i.e. specifically the dim tags.
      https://github.com/rwth-i6/returnn/issues/706
    Throws an exception if this is not the case.

    :param set[Dim|_MarkedDim]|tuple|list out_shape:
      It must be a set, with the only exception when it is empty (then it doesn't matter).
      See :func:`dim_tags_set`.
    """
    self_dim_tags = self.dim_tags_set_implicit
    self_dim_tags_implicit_only = self.dim_tags_set_implicit_only_wrapped
    if not out_shape:
      if self_dim_tags:
        raise VerifyOutShapeException(
          "%s verify_out_shape, with dims %s, does not match empty out_shape %r" % (self, self_dim_tags, out_shape))
      return
    if not isinstance(out_shape, set):
      raise TypeError("%s verify_out_shape: expects a set but got %s" % (self, type(out_shape)))
    remaining = set(self_dim_tags)
    for dim in out_shape:
      if isinstance(dim, Dim):
        dim_tag = dim
      elif isinstance(dim, _ImplicitDim):
        dim_tag = dim.tag
        if dim not in self_dim_tags_implicit_only:
          raise VerifyOutShapeException(
            "%s verify_out_shape, with dims %s, with out_shape %s, %s is not an implicit dim in self" % (
              self, self_dim_tags, out_shape, dim))
      elif isinstance(dim, OptionalDim):
        dim_tag = dim.tag
        if dim_tag not in remaining:
          continue
      else:
        raise TypeError("%s verify_out_shape with out_shape %s: expect dim tags but got %s" % (
          self, out_shape, type(dim)))
      if dim_tag not in remaining:
        if dim_tag in self_dim_tags:  # can happen e.g. if specified once as implicit dim and then also as explicit
          raise VerifyOutShapeException(
            "%s verify_out_shape, with dims %s, does not match out_shape %r, dim %s multiple times in out_shape" % (
              self, self_dim_tags, out_shape, dim))
        raise VerifyOutShapeException(
          "%s verify_out_shape, with dims %s, does not match out_shape %r, %s not in self" % (
            self, self_dim_tags, out_shape, dim))
      remaining.discard(dim_tag)
    if remaining:
      raise VerifyOutShapeException(
        "%s verify_out_shape, dims %s are not specified in out_shape %s" % (self, remaining, out_shape))

  def get_placeholder_kwargs(self, with_batch=True):
    """
    :param bool with_batch:
    :return: kwargs for tf.compat.v1.placeholder
    :rtype: dict[str]
    """
    return dict(name=self.name, dtype=self.dtype, shape=self.batch_shape if with_batch else self.shape)

  def get_axes_with_size(self):
    """
    :return: list of axes which can vary in size for each entry of the batch-dim, e.g. the time-dim-axis.
      The axis index is counted without the batch-dim.
    :rtype: list[int]
    """
    return [i for (i, dim) in enumerate(self.shape) if dim is None]

  def get_kwargs(self, include_special_axes=True):
    """
    :param bool include_special_axes: whether to include time and feature special axis marker
    :return: relevant attrib items for copying
    :rtype: dict[str]
    """
    keys = ["name", "dim_tags", "dtype"]
    if include_special_axes:
      keys += ["time_dim_axis"]
      if self._feature_dim_axis is not NotSpecified:
        keys += ["feature_dim_axis"]
    if self.sparse_dim:
      # Sparse is False by default. And the dim is inferred from the feature dim, or otherwise does not make sense.
      keys += ["sparse_dim"]
    if self.batch is not None:
      keys += ["batch"]
    if self.beam is not None:
      keys += ["beam"]
    if self.control_flow_ctx:
      keys += ["control_flow_ctx"]
    if not self.available_for_inference:
      keys += ["available_for_inference"]
    return {key: getattr(self, key) for key in keys}

  def get_description(self, with_name=True, with_placeholder=False, catch_exceptions=False):
    """
    :param bool with_name:
    :param bool with_placeholder:
    :param bool catch_exceptions:
    :return: description of self. also used for __repr__
    :rtype: str
    """
    # Avoid redundant information (most information is covered in batch_shape_meta).
    # Also try to avoid confusion (e.g. `shape` vs `batch_shape`).
    keys = []
    if self.sparse:
      keys.append("dtype")
      keys.append("sparse_dim")
    else:
      if self.dtype != "float32":
        keys.append("dtype")
    if with_placeholder:
      keys.append("placeholder")
    if not self.available_for_inference:
      keys.append("available_for_inference")
    if self.beam is not None:
      # With batch, it should be contained already in batch_shape_meta (if everything is correct).
      # We anyway add it, in case sth is incorrect or incomplete.
      if not self.batch or self.batch.beam != self.beam:
        keys.append("beam")
    args = []
    if with_name:
      name = getattr(self, "name", None)
      args += [repr(name) if name else "<undefined>"]
    try:
      batch_shape_meta = "[%s]" % ",".join(self.get_batch_axes_short_description())
    except Exception as exc:
      if catch_exceptions:
        batch_shape_meta = "<!%s: %s>" % (type(exc).__name__, exc)
      else:
        raise
    args += [batch_shape_meta]
    for key in keys:
      try:
        value_repr = repr(getattr(self, key))
      except Exception as exc:
        if catch_exceptions:
          value_repr = "<!%s: %s>" % (type(exc).__name__, exc)
        else:
          raise
      args += ["%s=%s" % (key, value_repr)]
    if self.control_flow_ctx:
      try:
        value_repr = self.control_flow_ctx.repr_inner()
      except Exception as exc:
        if catch_exceptions:
          value_repr = "<!%s: %s>" % (type(exc).__name__, exc)
        else:
          raise
      args += ["ctx=" + value_repr]
    return "Data{%s}" % ", ".join(args)

  def get_batch_axes_short_description(self, special_axes=True):
    """
    :param bool special_axes: special markers for old-style time_dim_axis and feature_dim_axis
    :rtype: list[str]
    """
    res = []
    for axis, dim_tag in enumerate(self.dim_tags):
      descriptions = []
      if axis == self.batch_dim_axis:
        if self.batch:
          descriptions.append(self.batch.short_repr())
        else:
          descriptions.append("B?")
      if special_axes:
        if axis == self.time_dim_axis:
          descriptions.append("T")
        if axis == self.feature_dim_axis:
          descriptions.append("F")
      if self.batch_shape[axis] is None:
        if axis == self.batch_dim_axis:
          pass  # expected
        else:
          descriptions.append(dim_tag.short_repr())
      elif axis != self.batch_dim_axis or not self.batch:
        descriptions.append(dim_tag.short_repr())
      res.append("|".join(descriptions) or "?")
    return res

  def get_compare_key(self):
    """
    :return: some key which can be used for compare functions, i.e. such that
      cmp(get_compare_key(self), get_compare_key(other)) == cmp(self, other),
      i.e. we define some order by that.
      Note that this order is not totally fixed, and might change.
    :rtype: object
    """
    return (
      self.dtype,
      self.shape,
      self.batch_dim_axis, self.feature_dim_axis, self.time_dim_axis,
      self.dim_tags,
      self.batch, self.beam)

  def __repr__(self):
    return self.get_description(catch_exceptions=True)

  def __hash__(self):
    return id(self)

  def __getstate__(self):
    d = vars(self)
    d["_batch"] = None  # BatchInfo pickling not supported
    d["_placeholder"] = None  # do not store the TF tensors
    return d

  def _adapt_batch_consistent_dim_tags(self):
    if not self.batch:  # uninitialized
      return
    self._dim_tags = tuple(
      tag.get_for_batch_ctx(batch=self.batch, ctx=self.control_flow_ctx) for tag in self._dim_tags)

  def copy(self, name=None):
    """
    :param str name: if given, will overwrite this name
    :return: copy of myself, using self.get_kwargs(), and with placeholder and size_placeholder
    :rtype: Data
    """
    data = Data(**self.get_kwargs())
    data.placeholder = self.placeholder
    if name:
      data.name = name
    return data

  def copy_as_batch_major(self):
    """
    :return: copy of myself with batch_dim_axis == 0
    :rtype: Data
    """
    return self.copy_with_batch_dim_axis(0)

  def copy_as_time_major(self):
    """
    :return: copy of myself with time_dim_axis == 0
    :rtype: Data
    """
    assert self.time_dim_axis is not None
    return self.copy_with_time_dim_axis(0)

  def copy_with_batch_dim_axis(self, batch_dim_axis):
    """
    :param int batch_dim_axis:
    :return: copy of myself with specific batch_dim_axis
    :rtype: Data
    """
    assert self.batch_dim_axis is not None
    return self.copy_move_axis(self.batch_dim_axis, batch_dim_axis)

  def copy_with_time_dim_axis(self, time_dim_axis):
    """
    :param int time_dim_axis:
    :return: copy of myself with specific time_dim_axis
    :rtype: Data
    """
    assert self.time_dim_axis is not None
    return self.copy_move_axis(self.time_dim_axis, time_dim_axis)

  def copy_transpose(self, perm):
    """
    :param list[int] perm: permutation of the axes, counted with batch-dim.
      Maps the new axes to the old axes
    :return: copy of myself with permuted axes
    :rtype: Data
    """
    assert len(perm) == self.batch_ndim
    assert set(perm) == set(range(self.batch_ndim))
    if all(perm[axis] == axis for axis in range(self.batch_ndim)):
      return self.copy()
    inv_perm_ = {j: i for (i, j) in enumerate(perm)}
    inv_perm = [j for (i, j) in sorted(inv_perm_.items())]

    def translate_axis(axis):
      """
      :param int|None axis: counted with batch-dim
      :return: translated axis (if not None)
      :rtype: int|None
      """
      if axis is None:
        return None
      return inv_perm[axis]

    data_opts = self.get_kwargs(include_special_axes=False)
    if self.placeholder is not None:
      from returnn.tf.util.basic import get_valid_scope_name_from_str
      data_opts["placeholder"] = tf.transpose(
        self.placeholder, perm, name="%s_transpose" % get_valid_scope_name_from_str(self.name))
    data_opts["dim_tags"] = tuple(self.dim_tags[perm[i]] for i in range(self.batch_ndim))
    data = Data(**data_opts)
    data.sanity_check()
    if self.time_dim_axis is not None and translate_axis(self.time_dim_axis) != data.time_dim_axis:
      data.time_dim_axis = translate_axis(self.time_dim_axis)
    if self.feature_dim_axis is not None and translate_axis(self.feature_dim_axis) != data.feature_dim_axis:
      data.feature_dim_axis = translate_axis(self.feature_dim_axis)
    return data

  def copy_move_axis(self, old_axis, new_axis):
    """
    :param int old_axis: counted with batch-dim
    :param int new_axis: counted with batch-dim
    :return: copy of myself with moved axis (see :func:`move_axis`)
    :rtype: Data
    """
    if old_axis < 0:
      old_axis += self.batch_ndim
      assert old_axis >= 0
    assert 0 <= old_axis < self.batch_ndim
    if new_axis < 0:
      new_axis += self.batch_ndim
      assert new_axis >= 0
    assert 0 <= new_axis < self.batch_ndim
    if old_axis == new_axis:
      return self.copy()

    perm = list(range(self.batch_ndim))
    old = perm.pop(old_axis)
    perm.insert(new_axis, old)
    return self.copy_transpose(perm)

  def copy_swap_axes(self, axis1, axis2):
    """
    Like :func:`Data.copy_move_axis`, but keeps all other axes unchanged.
    :param int axis1: counted with batch-dim
    :param int axis2: counted with batch-dim
    :return: copy of myself with moved axis (see :func:`swapaxes`)
    :rtype: Data
    """
    if axis1 < 0:
      axis1 += self.batch_ndim
    assert 0 <= axis1 < self.batch_ndim
    if axis2 < 0:
      axis2 += self.batch_ndim
    assert 0 <= axis2 < self.batch_ndim
    if axis1 == axis2:
      return self.copy()

    perm = list(range(self.batch_ndim))
    perm[axis1], perm[axis2] = perm[axis2], perm[axis1]
    return self.copy_transpose(perm)

  def copy_as_bt_or_tb_major(self):
    """
    :rtype: Data
    :return: copy of myself in batch-time-major or time-batch-major
    """
    assert self.have_batch_axis() and self.have_time_axis()
    if self.batch_dim_axis == 0:
      return self.copy_with_time_dim_axis(1)
    if self.time_dim_axis == 0:
      return self.copy_with_batch_dim_axis(1)
    if self.batch_dim_axis > self.time_dim_axis:
      return self.copy_as_time_major().copy_as_bt_or_tb_major()
    return self.copy_as_batch_major().copy_as_bt_or_tb_major()

  def copy_with_feature_dim_axis(self, feature_dim_axis):
    """
    :param int feature_dim_axis: can also be negative
    :return: copy of myself with specific feature dim axis
    :rtype: Data
    """
    assert self.feature_dim_axis is not None
    return self.copy_move_axis(self.feature_dim_axis, feature_dim_axis)

  def copy_as_batch_feature_major(self):
    """
    :return: copy of self with batch_dim_axis == 0 and feature_dim_axis == 1
    :rtype: Data
    """
    assert self.batch_dim_axis is not None
    assert self.feature_dim_axis is not None
    data = self.copy_as_batch_major()
    data = data.copy_with_feature_dim_axis(1)
    return data

  def copy_as_time_batch_major(self):
    """
    :return: copy of self with batch_dim_axis == 1 and time_dim_axis == 0
    :rtype: Data
    """
    assert self.have_batch_axis() and self.have_time_axis()
    data = self.copy_as_bt_or_tb_major()
    if data.time_dim_axis == 1:
      data = data.copy_move_axis(0, 1)
    return data

  def copy_as_batch_spatial_major(self):
    """
    :return: copy with batch_dim_axis == 0, then all dynamic axes, then any other spatial axes, last feature axis
    :rtype: Data
    """
    data = self.copy_as_batch_major()
    if data.feature_dim_axis is not None:
      data = data.copy_with_feature_last()
    if data.size_placeholder:
      for i, (j, size) in enumerate(sorted(data.size_placeholder.items())):
        data = data.copy_move_axis(data.get_batch_axis(j), i + 1)
    if data.feature_dim_axis is not None:
      assert data.feature_dim_axis == data.batch_ndim - 1
      # Maybe reset feature_dim_axis to unspecified.
      if data.feature_dim_axis_or_unspecified is not NotSpecified:
        if data._default_feature_dim_axis() == data.feature_dim_axis:
          data.feature_dim_axis = NotSpecified
    return data

  def copy_with_feature_last(self):
    """
    :return: copy of self with feature_dim_axis being the very last axis
    :rtype: Data
    """
    assert self.feature_dim_axis is not None
    return self.copy_with_feature_dim_axis(-1)

  def copy_add_batch_dim(self, batch_dim_axis, batch=None, dim_tag=None):
    """
    :param int batch_dim_axis:
    :param BatchInfo|None batch:
    :param Dim|None dim_tag:
    :return: copy of myself with added batch-dim
    :rtype: Data
    """
    assert self.batch_dim_axis is None
    if not batch:
      from returnn.tf.layers.base import LayerBase
      batch = LayerBase.get_recent_layer().get_batch_info()
    if batch_dim_axis < 0:
      assert batch_dim_axis + self.batch_ndim + 1 >= 0
      batch_dim_axis += self.batch_ndim + 1
    assert 0 <= batch_dim_axis <= self.batch_ndim
    data_opts = self.get_kwargs(include_special_axes=False)
    placeholder = self.placeholder
    if placeholder is not None:
      from .basic import get_valid_scope_name_from_str
      placeholder = tf.expand_dims(
        self.placeholder, batch_dim_axis, name=get_valid_scope_name_from_str("%s_add_batch_dim" % self.name))
      if not isinstance(batch.dim, int) or batch.dim != 1:
        tiles = [1] * batch_dim_axis + [batch.dim] + [1] * (self.batch_ndim - batch_dim_axis)
        placeholder = tf.tile(placeholder, tiles)
    dim_tags = list(self.dim_tags)
    if dim_tag:
      assert dim_tag.is_batch_dim()
      assert dim_tag.dimension == batch.static_dim or dim_tag.dimension is None
      assert dim_tag.batch == batch
    else:
      dim_tag = Dim(
        kind=Dim.Types.Batch, description="batch", dimension=batch.static_dim, batch=batch)
    dim_tags.insert(batch_dim_axis, dim_tag)
    data_opts["dim_tags"] = dim_tags
    data_opts["batch"] = batch
    data_opts["beam"] = batch.beam
    other_special_axes = self.get_special_axes_dict(counted_with_batch_dim=True, only_available=True)
    for k, a in other_special_axes.items():
      data_opts[k] = a if (a < batch_dim_axis) else (a + 1)
    return Data(placeholder=placeholder, **data_opts)

  def copy_add_spatial_dim(self, spatial_dim_axis=None, dim=1, auto_time_dim_axis=True):
    """
    :param int|None spatial_dim_axis: counted with batch-dim. if there is no time-dim, this will be it.
    :param int|None dim:
    :param bool auto_time_dim_axis:
    :return: copy of myself with added spatial-dim
    :rtype: Data
    """
    if dim is None:
      assert not self.placeholder
    dim_tag = SpatialDim("added_spatial", dimension=dim)
    if spatial_dim_axis is None:
      spatial_dim_axis = self.get_default_new_axis_for_dim_tag(dim_tag)
    v = self.copy_add_dim_by_tag(dim_tag, unbroadcast=True, axis=spatial_dim_axis)
    if auto_time_dim_axis and self.time_dim_axis is None:
      v.time_dim_axis = spatial_dim_axis
    return v

  def copy_add_feature_dim(self, axis=None):
    """
    :param int|None axis:
    :return: self with a new feature dim axis with dim 1.
      If there is an existing feature dim, the new feature dim will be added right after.
      If we are sparse, we don't add a feature dim, but it becomes a spatial dim instead.
    :rtype: Data
    """
    if self.sparse:
      # By definition, we don't have a feature dim. We allow this though. We just make it a spatial axis.
      return self.copy_add_spatial_dim(spatial_dim_axis=axis)
    dim_tag = FeatureDim("feature1", dimension=1)
    if axis is None:
      axis = self.get_default_new_axis_for_dim_tag(dim_tag)
    v = self.copy_add_dim_by_tag(dim_tag, axis=axis)
    if v.feature_dim_axis_or_unspecified is not NotSpecified:
      v.feature_dim_axis = NotSpecified
    if axis < 0:
      axis += v.batch_ndim
      assert axis >= 0
    assert 0 <= axis < v.batch_ndim
    if v.feature_dim_axis != axis:
      v.feature_dim_axis = axis
    return v

  def get_default_new_axis_for_dim_tag(self, dim_tag):
    """
    :param Dim dim_tag:
    :rtype: int
    """
    if dim_tag.is_batch_dim():
      return 0
    # Note: if dim_tag is feature, but we are sparse, we just treat is as spatial, handled below.
    if dim_tag.is_feature_dim() and not self.sparse:
      if self.feature_dim_axis is not None:
        return self.feature_dim_axis + 1  # after existing feature-dim
      else:
        return self.batch_ndim  # at the end
    assert dim_tag.is_spatial_dim() or (dim_tag.is_feature_dim() and self.sparse)
    if dim_tag.dimension is None and self.get_dynamic_axes():
      return self.get_dynamic_axes()[-1] + 1  # after existing dynamic axis
    if self.get_spatial_batch_axes():
      return self.get_spatial_batch_axes()[-1] + 1  # after the existing spatial dim
    elif self.feature_dim_axis is not None:
      return self.feature_dim_axis  # add it before the feature dim
    else:
      return self.batch_ndim  # add it at the end

  def copy_add_dim_by_tag(self, dim_tag, unbroadcast=False, axis=None):
    """
    :param Dim dim_tag:
    :param bool unbroadcast: If True unbroadcast the newly added axis.
      Will infer the unbroadcast shape via :func:`Dim.get_dim_value`
    :param int|None axis:
    :rtype: Data
    """
    assert dim_tag.can_be_used_as_dim()
    from .basic import get_valid_scope_name_from_str
    if axis is None:
      axis = self.get_default_new_axis_for_dim_tag(dim_tag=dim_tag)
    if axis < 0:
      axis += self.batch_ndim + 1
    assert 0 <= axis <= self.batch_ndim

    if dim_tag.is_batch_dim():
      if unbroadcast:
        return self.copy_add_batch_dim(batch_dim_axis=axis, batch=dim_tag.batch, dim_tag=dim_tag)
      else:
        batch_info = BatchInfo.make_global_broadcast_batch_info()
        return self.copy_add_batch_dim(
          batch_dim_axis=axis, batch=batch_info,
          dim_tag=dim_tag if (dim_tag.dimension == 1 and dim_tag.batch == batch_info) else None)

    data_opts = self.get_kwargs()
    # Note: if dim_tag is feature, but we are sparse, we just make it spatial
    if self.sparse and dim_tag.is_feature_dim():
      dim_tag = dim_tag.copy(same_as_self=True, kind=Dim.Types.Spatial)
    if not unbroadcast and dim_tag.dimension != 1:
      dim_tag = Dim(
        kind=dim_tag.kind, description="%s_dummy_dim1" % (dim_tag.description or "unnamed"), dimension=1,
        auto_generated=True)
    data_opts["dim_tags"] = self.dim_tags[:axis] + (dim_tag,) + self.dim_tags[axis:]
    other_special_axes = self.get_special_axes_dict(counted_with_batch_dim=True, only_available=True)
    for k, a in other_special_axes.items():
      data_opts[k] = a if (a < axis) else (a + 1)
    if dim_tag.is_feature_dim() and self.feature_dim_axis is None:
      data_opts.pop("feature_dim_axis", None)  # fall back to default
    if dim_tag.is_spatial_dim() and self.time_dim_axis is None:
      data_opts.pop("time_dim_axis", None)  # fall back to default
    if self.placeholder is not None:
      with tf.name_scope("%s_copy_add_dim_by_tag" % get_valid_scope_name_from_str(self.name)):
        placeholder = tf.expand_dims(self.placeholder, axis)
        if dim_tag.dimension is None or dim_tag.dimension > 1:
          tiles = [1] * axis + [dim_tag.get_dim_value()] + [1] * (self.batch_ndim - axis)
          placeholder = tf.tile(placeholder, tiles)
      data_opts["placeholder"] = placeholder
    return Data(**data_opts)

  def copy_split_feature_dim(self, new_feature_dim):
    """
    Split it into (new_spatial_dim, new_feat_dim), in that order.
    This will increase the feature_dim_axis by one.

    :param int new_feature_dim: will be the new dim
    :rtype: Data
    """
    from .basic import get_shape
    assert not self.sparse
    assert self.feature_dim_axis is not None
    assert self.dim is not None
    assert self.dim % new_feature_dim == 0, "must be a multiple of the input feature dim"
    feature_dim_rem = self.dim // new_feature_dim
    new_feature_dim_axis = self.feature_dim_axis + 1
    data_opts = self.get_kwargs(include_special_axes=False)
    dim_tag_split_rem = Dim(
      kind=Dim.Types.Spatial, description="feature_split_rem_%i" % feature_dim_rem, auto_generated=True,
      dimension=feature_dim_rem)
    dim_tag_new = Dim(
      kind=self.dim_tags[self.feature_dim_axis].kind,
      description="feature_split_new_%i" % new_feature_dim, auto_generated=True,
      dimension=new_feature_dim)
    dim_tags = (
      self.dim_tags[:self.feature_dim_axis] +
      (dim_tag_split_rem, dim_tag_new) +
      self.dim_tags[self.feature_dim_axis + 1:])
    data_opts["dim_tags"] = dim_tags
    other_special_axes = self.get_special_axes_dict(counted_with_batch_dim=True, only_available=True)
    other_special_axes.pop("feature_dim_axis", None)
    for k, a in other_special_axes.items():
      data_opts[k] = a if (a < new_feature_dim_axis) else (a + 1)
    if self.placeholder is not None:
      self.placeholder.set_shape(self.batch_shape)
      old_shape = get_shape(self.placeholder)
      new_shape = (
        old_shape[:self.feature_dim_axis] +
        [feature_dim_rem, new_feature_dim] +
        old_shape[self.feature_dim_axis + 1:])
      data_opts["placeholder"] = tf.reshape(self.placeholder, new_shape, name="copy_split_feature_dim")
    return Data(**data_opts)

  def copy_extend_batch(self, batch):
    """
    Similar as copy_compatible_to with unbroadcast=True,
    we would possibly extend/expand our batch dim.
    See :class:`BatchInfo`.
    This assumes that we already have a batch dim
    (otherwise see :func:`copy_add_batch_dim`).

    This excludes any beam expansion, which is handled explicitly elsewhere
    (e.g. see :func:`copy_extend_with_beam`).

    :param BatchInfo batch:
    :rtype: Data
    """
    assert self.have_batch_axis()
    assert self.batch, "%s: batch unset" % self
    data = self.copy()
    batch = batch.copy_set_beam(data.beam)
    if data.batch.beam != data.beam:  # Check for some buggy code.
      data.batch = data.batch.copy_set_beam(data.beam)
    if data.batch == batch:
      return data
    data.batch = batch
    self._adapt_batch_consistent_dim_tags()
    if self.placeholder is not None:
      # This can only work if the batch is expanded.
      assert set(self.batch.virtual_dims).issubset(batch.virtual_dims)
      from .basic import get_shape
      from returnn.util.basic import ensure_list_of_type
      with tf.name_scope("copy_extend_batch"):
        axis = self.batch_dim_axis
        x = self.placeholder
        shape = get_shape(x)
        # Only fixed dims supported/implemented (no packed dims).
        old_dims = ensure_list_of_type(self.batch.virtual_dims, BatchInfo.FixedDim)
        new_dims = ensure_list_of_type(batch.virtual_dims, BatchInfo.FixedDim)
        batch_broadcast_shape = []  # type: typing.List[typing.Union[tf.Tensor,int]]  # fill below
        ndim_batch_split = self.batch_ndim - 1 + len(new_dims)
        tiles = [1] * ndim_batch_split  # type: typing.List[typing.Union[tf.Tensor,int]]  # overwrite below
        old_idx = 0
        for new_idx, new_dim in enumerate(new_dims):
          old_dim = old_dims[old_idx] if old_idx < len(old_dims) else None
          if old_dim == new_dim:
            batch_broadcast_shape.append(old_dim.size)
            old_idx += 1
          else:
            batch_broadcast_shape.append(1)
            tiles[axis + new_idx] = new_dim.size
        assert old_idx == len(old_dims)
        shape_batch_split = shape[:axis] + batch_broadcast_shape + shape[axis + 1:]
        x = tf.reshape(x, shape_batch_split)
        x = tf.tile(x, tiles)
        shape = shape[:axis] + [batch.dim] + shape[axis + 1:]
        x = tf.reshape(x, shape)
        data.placeholder = x
    return data

  def copy_compatible_to(self, data, add_dims=True, unbroadcast=False, except_feature=False, except_axis=None,
                         check_sparse=True, check_dtype=True):
    """
    :param Data data: other data which the returned tensor should be compatible to
      It would add any missing axes with a dim 1 axis for automatic broadcasting (with add_dims=True).
      It currently does not check whether existing dims match.
    :param bool add_dims: whether to add (broadcast, or unbroadcasted) dims. throws error if missing dim
    :param bool unbroadcast: if True, all added broadcast axes (axes with dim 1) will be tiled such that they match
    :param bool except_feature: if unbroadcast, do not unbroadcast the feature dim
    :param Dim|int|None except_axis: if unbroadcast, do not unbroadcast this axis
    :param bool check_sparse:
    :param bool check_dtype:
    :returns: Data, might add broadcast dimensions
    :rtype: Data
    """
    assert not check_sparse or self.sparse == data.sparse
    assert not check_dtype or self.dtype == data.dtype
    v = self.copy()
    if v.have_batch_axis() and data.have_batch_axis() and v.batch and data.batch and v.batch != data.batch:
      v = v.copy_extend_batch(data.batch)
    v.sparse_dim = data.sparse_dim  # we will later reset it. this is to better count the axes (feature and spatial)
    if v.batch_dim_axis is not None and data.batch_dim_axis is None:
      raise ValueError("copy_compatible_to: self %r has batch-dim, but target data %r has not" % (self, data))
    if data.batch_ndim < v.batch_ndim:
      raise ValueError("copy_compatible_to: self %r already has more dims than target data %r" % (self, data))

    is_equal_opts = dict(
      allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True, ignore_feature_dim=True)
    mapped_axes = data.find_matching_dim_map(v, list(range(v.batch_ndim)), is_equal_opts)  # maps v -> data
    assert len(mapped_axes) == v.batch_ndim

    except_axis_int = data.get_axis_from_description(except_axis, allow_int=True) if except_axis is not None else None

    for target_axis in range(data.batch_ndim):
      new_v_axis = min(target_axis, v.batch_ndim)
      if target_axis not in mapped_axes.values():
        if not add_dims:
          raise ValueError(
            "%s.copy_compatible_to(%s) not allowed, axis %i (%s) not in source" % (
              self, data, target_axis, data.dim_tags[target_axis]))
        # Dim in data, but not in v
        unbroadcast_axis = unbroadcast and not (
          except_feature and data.feature_dim_axis == target_axis) and not (
          except_axis_int is not None and except_axis_int == target_axis)
        v = v.copy_add_dim_by_tag(data.get_dim_tag(target_axis), axis=new_v_axis, unbroadcast=unbroadcast_axis)
        # Keep mapped_axes consistent
        mapped_axes = {v_ax + (1 if v_ax >= new_v_axis else 0): trg_ax for v_ax, trg_ax in mapped_axes.items()}
        mapped_axes[new_v_axis] = target_axis
      else:
        # Dim exists in both data and in v. Maybe order is wrong.
        matching_v_axes = [v_ax for v_ax, trg_ax in mapped_axes.items() if trg_ax == target_axis]
        assert len(matching_v_axes) == 1
        matching_v_axis = matching_v_axes[0]
        if target_axis != matching_v_axis:
          # Order was wrong
          v = v.copy_swap_axes(matching_v_axis, new_v_axis)
          # Keep mapped_axes consistent
          mapped_axes[matching_v_axis], mapped_axes[new_v_axis] = mapped_axes[new_v_axis], mapped_axes[matching_v_axis]

    assert v.batch_ndim == data.batch_ndim
    assert all(mapped_axes[ax] == ax for ax in range(v.batch_ndim))

    # Ensure time_dim_axis and feature_dim_axis is same as in data
    assert v.batch_dim_axis == data.batch_dim_axis  # there is only at most one batch_dim_axis
    v.time_dim_axis = data.time_dim_axis
    v.feature_dim_axis = data.feature_dim_axis_or_unspecified

    # Reset sparse
    if self.sparse:
      v.feature_dim_axis = NotSpecified
      v.sparse_dim = self.sparse_dim  # reset

    v.sanity_check()
    return v

  def copy_time_flattened(self):
    """
    :return: copy of myself where the time-axis is flattened away into the batch-dim-axis.
      See :func:`get_placeholder_time_flattened` and :func:`flatten_with_seq_len_mask for more details.
    :rtype: Data
    """
    assert self.batch_dim_axis is not None
    assert self.time_dim_axis is not None
    data_opts = self.get_kwargs(include_special_axes=False)
    if self.placeholder is not None:
      data_opts["placeholder"] = self.get_placeholder_time_flattened()
    dim_tag = self.dim_tags[self.time_dim_axis]
    dim_tag = Dim(
      kind=Dim.Types.Spatial, description="%s_flattened" % (dim_tag.description or "unnamed"), auto_generated=True)
    data_opts["dim_tags"] = (
      (dim_tag,) +
      tuple(tag for (i, tag) in enumerate(self.dim_tags) if i not in (self.batch_dim_axis, self.time_dim_axis)))
    data_opts["time_dim_axis"] = None
    data_opts.pop("feature_dim_axis", None)
    return Data(**data_opts)

  def copy_extend_with_beam(self, beam):
    """
    :param SearchBeam|None beam:
    :return: copy of myself where the batch-dim is extended/multiplied by beam_size, using tile_transposed
    :rtype: Data
    """
    from .basic import get_valid_scope_name_from_str, same_control_flow_ctx, tile_transposed
    data = self.copy()
    if data.beam and data.beam == beam:
      return data
    assert data.beam is None, "incompatible beam (%r vs %r)" % (data.beam, beam)
    if beam is None:
      return data
    data.beam = beam
    assert data.batch
    data.batch = data.batch.copy_set_beam(beam)
    with tf.name_scope("%s_data_extend_with_beam" % get_valid_scope_name_from_str(self.name)):
      if data.placeholder is not None:
        with same_control_flow_ctx(data.placeholder):
          data.placeholder = tile_transposed(data.placeholder, axis=data.batch_dim_axis, multiples=beam.beam_size)
          setattr(data.placeholder, "_RETURNN_beam_expanded_base_data", self)
    data._adapt_batch_consistent_dim_tags()
    return data

  def copy_merge_into_batch(self, axes):
    """
    :param list[int] axes: All axes to be merged into the batch axis.
      Must include the batch_dim_axis. The order is kept.
    :return: copy of myself where the the given axes are merged into the batch dim
    :rtype: Data
    """
    assert self.batch
    assert self.batch_dim_axis in axes
    assert sorted(set(axes)) == sorted(axes)
    min_axis = min(axes)
    axes = list(axes)
    data = self.copy()
    if axes != list(range(min_axis, min_axis + len(axes))):
      rem_axes_start = list(range(min_axis))
      rem_axes_end = [a for a in range(min_axis, self.batch_ndim) if a not in axes]
      data = data.copy_transpose(rem_axes_start + axes + rem_axes_end)
      axes = list(range(min_axis, min_axis + len(axes)))
      assert data.batch_dim_axis in axes
    tensor = data.placeholder
    batch = data.batch
    data = data.copy_template()
    batch_idx = 0
    for axis in axes:
      if axis == data.batch_dim_axis:
        batch_idx = len(batch.virtual_dims)  # add all remaining axes behind
        continue
      batch = batch.copy_extend_with_padded_or_fixed_dim_tag(
        dim_tag=data.dim_tags[axis], new_dim_idx=batch_idx)
      batch_idx += 1
    for axis in reversed(sorted(axes)):
      if axis != data.batch_dim_axis:
        data = data.copy_template_excluding_axis(axis)
    data.batch = batch
    if tensor is not None:
      from .basic import get_shape
      shape = get_shape(tensor)
      tensor = tf.reshape(tensor, shape[:min_axis] + [-1] + shape[min_axis + len(axes):])
      data.placeholder = tensor
    return data

  def copy_squeeze_axes(self, axes):
    """
    :param list[int] axes: counted with batch dim
    :return: copy of myself, with squeezed axes
    :rtype: Data
    """
    from .basic import get_valid_scope_name_from_str
    assert isinstance(axes, (list, tuple))
    assert all(self.batch_shape[axis] == 1 for axis in axes)
    assert all(0 <= axis < self.batch_ndim for axis in axes)
    if not axes:
      return self.copy()
    data_opts = self.get_kwargs(include_special_axes=False)
    if self.placeholder is not None:
      data_opts["placeholder"] = tf.squeeze(
        self.placeholder, axes,
        name="%s_squeeze_axes" % get_valid_scope_name_from_str(self.name))
    data_opts["dim_tags"] = [tag for (i, tag) in enumerate(self.dim_tags) if i not in axes]
    if self.time_dim_axis is not None:
      if self.time_dim_axis in axes:
        data_opts.pop("time_dim_axis", None)
      else:
        data_opts["time_dim_axis"] = self.time_dim_axis - len([axis for axis in axes if axis < self.time_dim_axis])
    if not self.sparse:
      if self.feature_dim_axis is not None and self.feature_dim_axis_or_unspecified is not NotSpecified:
        if self.feature_dim_axis in axes:
          data_opts.pop("feature_dim_axis", None)
        else:
          data_opts["feature_dim_axis"] = (
            self.feature_dim_axis - len([axis for axis in axes if axis < self.feature_dim_axis]))
    return Data(**data_opts)

  def copy_template(self, name=None, dtype=None):
    """
    :param str|None name:
    :param str|None dtype:
    :return: copy of myself, using self.get_kwargs(), without placeholder
    :rtype: Data
    """
    kwargs = self.get_kwargs()
    if name:
      kwargs["name"] = name
    if dtype:
      kwargs["dtype"] = dtype
    return Data(**kwargs)

  def copy_template_dense(self, name=None, dtype=None):
    """
    :param str|None name:
    :param str|None dtype:
    :return: copy of myself, using self.get_kwargs(), without placeholder
    :rtype: Data
    """
    out = self.copy_template(name=name)
    if out.sparse:
      feat_dim = out.sparse_dim
      out.sparse = False
      out.dtype = "float32"
      out = out.copy_add_dim_by_tag(dim_tag=feat_dim, unbroadcast=True, axis=-1)
      out.feature_dim_axis = NotSpecified
      assert out.feature_dim_axis == out.batch_ndim - 1
    if dtype:
      out.dtype = dtype
    return out

  def copy_template_excluding_axis(self, exclude_axis, name=None):
    """
    :param int exclude_axis: axis to be removed.
    :param str|None name: if set, this will be the new name.
    :return: copy of myself excluding exclude_axis axis, without placeholder.
    :rtype: Data
    """
    kwargs = self.get_kwargs(include_special_axes=False)
    if exclude_axis < 0:
      exclude_axis += self.batch_ndim
      assert exclude_axis >= 0
    assert 0 <= exclude_axis < self.batch_ndim
    if exclude_axis == self.feature_dim_axis:
      kwargs.pop("dim", None)
    other_special_axes = self.get_special_axes_dict(counted_with_batch_dim=True, only_available=True)
    for axis_name, axis in other_special_axes.items():
      if axis == exclude_axis:
        continue
      kwargs[axis_name] = axis if (axis < exclude_axis) else (axis - 1)
    new_dim_tags = list(self.dim_tags)
    del new_dim_tags[exclude_axis]
    kwargs["dim_tags"] = new_dim_tags
    if name:
      kwargs["name"] = name
    return Data(**kwargs)

  def copy_template_excluding_spatial_dim(self, spatial_axis_num, name=None):
    """
    :param int spatial_axis_num: index in self.get_spatial_batch_axes()
    :param str|None name: if set, this will be the new name
    :return: copy of myself excluding the time-dimension without placeholder
    :rtype: Data
    """
    spatial_axes = self.get_spatial_batch_axes()
    if spatial_axis_num < 0:
      spatial_axis_num += len(spatial_axes)
      assert spatial_axis_num >= 0
    assert 0 <= spatial_axis_num < len(spatial_axes)
    axis_to_exclude = spatial_axes[spatial_axis_num]
    return self.copy_template_excluding_axis(exclude_axis=axis_to_exclude, name=name)

  def copy_template_excluding_time_dim(self, name=None):
    """
    :param str|None name: if set, this will be the new name
    :return: copy of myself excluding the time-dimension without placeholder
    :rtype: Data
    """
    assert self.time_dim_axis is not None
    return self.copy_template_excluding_axis(exclude_axis=self.time_dim_axis, name=name)

  def copy_template_adding_time_dim(self, name=None, time_dim_axis=0):
    """
    Adds a time-dim-axis.
    If a time-dim-axis already exists, it will anyway create this new one.

    :param str|None name: if set, this will be the new name
    :param int time_dim_axis: the new time-dim-axis index
    :return: copy of myself adding the time-dimension without placeholder
    :rtype: Data
    """
    if time_dim_axis < 0:
      time_dim_axis += self.batch_ndim + 1
      assert time_dim_axis >= 0
    assert 0 <= time_dim_axis <= self.batch_ndim
    kwargs = self.get_kwargs(include_special_axes=False)
    dim_tag = Dim(kind=Dim.Types.Time, description="unknown_time", dimension=None, auto_generated=True)
    dim_tags = self.dim_tags[:time_dim_axis] + (dim_tag,) + self.dim_tags[time_dim_axis:]
    kwargs["dim_tags"] = dim_tags
    other_special_axes = self.get_special_axes_dict(counted_with_batch_dim=True, only_available=True)
    other_special_axes.pop("time_dim_axis", None)
    for axis_name, axis in other_special_axes.items():
      kwargs[axis_name] = axis if (axis < time_dim_axis) else (axis + 1)
    kwargs["time_dim_axis"] = time_dim_axis
    if name:
      kwargs["name"] = name
    return Data(**kwargs)

  def copy_template_replace_dim_tag(self, axis, new_dim_tag, name=None):
    """
    :param int axis:
    :param Dim new_dim_tag:
    :param str|None name: new name
    :rtype: Data
    """
    assert new_dim_tag.can_be_used_as_dim()
    if axis < 0:
      assert axis + self.batch_ndim >= 0
      axis += self.batch_ndim
    assert 0 <= axis < self.batch_ndim
    opts = self.get_kwargs()
    if self.dim_tags[axis].is_batch_dim():
      opts.pop("batch", None)
    if new_dim_tag.is_batch_dim():
      if self.time_dim_axis == axis:
        opts.pop("time_dim_axis", None)
      if self.feature_dim_axis == axis:
        opts.pop("feature_dim_axis", None)
    dim_tags = self.dim_tags[:axis] + (new_dim_tag,) + self.dim_tags[axis + 1:]
    opts["dim_tags"] = dim_tags
    if self.feature_dim_axis_or_unspecified is not NotSpecified:
      if self.feature_dim_axis == axis and self.dim_tags[axis].is_feature_dim() and not new_dim_tag.is_feature_dim():
        opts["feature_dim_axis"] = None
    if name:
      opts["name"] = name
    return Data(**opts)

  def copy_template_replace_dim(self, axis, new_dim, new_size=None):
    """
    :param int axis:
    :param int|None new_dim:
    :param tf.Tensor|None new_size:
    :rtype: Data
    """
    dim_tag = self.dim_tags[axis]
    if dim_tag.is_batch_dim():
      assert new_dim is None
      return self.copy_template()  # nothing to do
    dim_tag = Dim(
      kind=dim_tag.kind, description="%s_replaced" % (dim_tag.description or "unnamed"),
      auto_generated=True,
      dimension=new_dim, dyn_size=new_size)
    return self.copy_template_replace_dim_tag(axis=axis, new_dim_tag=dim_tag)

  def copy_template_new_dim_tags(self, new_dim_tags, name=None, keep_special_axes=False):
    """
    :param list[Dim]|tuple[Dim] new_dim_tags:
    :param str|None name:
    :param bool keep_special_axes:
    :rtype: Data
    """
    if keep_special_axes:
      assert len(new_dim_tags) == self.batch_ndim
    opts = self.get_kwargs(include_special_axes=keep_special_axes)
    opts["dim_tags"] = new_dim_tags
    if name:
      opts["name"] = name
    return Data(**opts)

  def copy_template_set_ctx(self, ctx):
    """
    :param ControlFlowContext ctx:
    :return: new Data instance
    :rtype: Data
    """
    kwargs = self.get_kwargs()
    kwargs["control_flow_ctx"] = ctx
    return Data(**kwargs)

  def copy_template_unpack_batch(self):
    """
    If the batch dim contains a :class:`BatchInfo.PackedDim`, unpack it and restore the data from before the packing.

    :rtype: Data
    """
    assert self.have_batch_axis()
    assert self.batch, "%s: batch unset" % self
    data = self.copy()
    kwargs = self.get_kwargs(include_special_axes=False)

    dim_tags = []
    for dim_tag in data.dim_tags:
      if dim_tag.is_batch_dim() and dim_tag.batch and len(dim_tag.batch.virtual_dims) > 0:
        batch = dim_tag.batch
        new_batch_dim_tag = None
        for virtual_dim in batch.virtual_dims:
          if isinstance(virtual_dim, BatchInfo.PackedDim):
            dim_tags.append(virtual_dim.dim_tag)
            batch = batch.copy_remove_dim(virtual_dim)
          elif isinstance(virtual_dim, BatchInfo.GlobalBatchDim):
            assert not new_batch_dim_tag
            new_batch_dim_tag = Dim(kind=Dim.Types.Batch, description=dim_tag.description)
            dim_tags.append(new_batch_dim_tag)
        assert new_batch_dim_tag, "%s: batch info %r invalid" % (self, batch)
        new_batch_dim_tag.batch = batch
        kwargs["batch"] = batch
      else:
        dim_tags.append(dim_tag)

    kwargs["dim_tags"] = dim_tags
    return Data(**kwargs)

  def _get_variable_dim_pattern(self):
    """
    :return: tuple with bools specifying which dims of the shape (excluding batch-dim) are of variable length.
     e.g. (time,feature), shape=(None,128), this returns (True, False)
    :rtype: tuple[bool]
    """
    return tuple([dim is None for dim in self.shape])

  def _get_var_len_axes(self):
    return [i for (i, d) in enumerate(self._get_variable_dim_pattern()) if d]

  def matches_var_dim_pattern(self, other):
    """
    :param Data other:
    :return: whether the variable-dims pattern matches,
      i.e. same variable dims (get_variable_dim_pattern), same time dim, excluding batch-dim.
      i.e. the size_placeholder should be compatible.
    :rtype: bool
    """
    if self.time_dim_axis_excluding_batch != other.time_dim_axis_excluding_batch:
      return False
    return self._get_var_len_axes() == other._get_var_len_axes()

  @property
  def dim_tags(self):
    """
    :rtype: tuple[Dim]
    """
    return self._dim_tags

  @property
  def shape(self):
    """
    :return: shape without batch-dim. e.g. (time,feat) = (None,128)
    :rtype: tuple[int|None]
    """
    return tuple(tag.dimension for tag in self._dim_tags if not tag.is_batch_dim())

  @shape.setter
  def shape(self, shape):
    """
    :param tuple[int|None] shape:
    """
    if tuple(shape) == self.shape:
      return
    raise Exception("%s: setting the shape is not allowed (new shape %s)" % (self, shape))

  @property
  def batch_shape(self):
    """
    :return: shape with added batch-dim. e.g. (batch,time,feat) = (None,None,128)
    :rtype: tuple[int|None]
    """
    return tuple(tag.dimension for tag in self.dim_tags)

  # noinspection PyShadowingNames
  def get_batch_shape(self, batch_dim):
    """
    :param int|tf.Tensor|None batch_dim:
    :return: shape with added batch-dim. e.g. (batch,time,feat) = (None,None,128)
    :rtype: tuple[int|None]
    """
    if self.batch_dim_axis is not None:
      return self.shape[:self.batch_dim_axis] + (batch_dim,) + self.shape[self.batch_dim_axis:]
    return self.shape

  def get_dynamic_batch_shape(self):
    """
    :rtype: list[int|tf.Tensor]
    """
    return [self.get_dim(axis) for axis in range(self.batch_ndim)]

  def have_varying_shape_in_ctx(self):
    """
    :return: whether the (dynamic) shape can change in this control flow context.
      E.g. when self.control_flow_context is a loop, and we have one dynamic dim
      where dyn_size_ext has the same control_flow_context
      (such that dyn_size_ext has e.g. shape [B,T] outside the loop).
      This can be relevant for accumulating values of self.placeholder
      e.g. via tf.TensorArray.
    :rtype: bool
    """
    return any(tag.control_flow_ctx for tag in self.dim_tags)

  @property
  def size_placeholder(self):
    """
    For compatibility, return a proxy object which behaves like the original dict.

    :rtype: _SizePlaceholderProxy
    """
    return _SizePlaceholderProxy(self)

  @size_placeholder.setter
  def size_placeholder(self, sizes):
    """
    :param dict[int,tf.Tensor]|None sizes:
    """
    if sizes is None:
      return
    for axis_wo_b, size in sizes.items():
      self.set_dynamic_size(axis=self.get_batch_axis(axis_wo_b), sizes=size)

  @property
  def shape_dense(self):
    """
    :return: shape with feature dim axis
    :rtype: tuple[int|None]
    """
    if self.sparse:
      return self.shape + (self.dim,)  # by default, assume at the end
    return self.shape

  @property
  def batch_shape_dense(self):
    """
    :rtype: tuple[int|None]
    """
    if self.sparse:
      return self.batch_shape + (self.dim,)
    return self.batch_shape

  @property
  def dim_tags_sparse(self):
    """
    :return: dim tags without feature dim axis
    :rtype: tuple[Dim]
    """
    if self.sparse or not self.have_feature_axis():
      return self.dim_tags
    return self.dim_tags[:self.feature_dim_axis] + self.dim_tags[self.feature_dim_axis + 1:]

  @property
  def dim_tags_set_implicit_only_wrapped(self):
    """
    :return: Dim tags implicit by sparse dim, or dynamic sizes, and not present as explicit dims.
      Also see :func:`dim_tags_set`.
    :rtype: set[_ImplicitDim]
    """
    self_dim_tags = set(self.dim_tags)
    dims = set()
    if self.sparse_dim and self.sparse_dim not in self_dim_tags:
      dims.add(ImplicitSparseDim(self.sparse_dim))
    for dim in self.dim_tags:
      if dim.dyn_size_ext:
        for dim_ in dim.dyn_size_ext.dim_tags:
          if dim_ not in self_dim_tags:
            dims.add(ImplicitDynSizeDim(dim_))
    return dims

  @property
  def dim_tags_set_implicit_only(self):
    """
    :return: Dim tags implicit by sparse dim, or dynamic sizes, and not present as explicit dims.
      Also see :func:`dim_tags_set`.
    :rtype: set[Dim]
    """
    return set(dim.tag for dim in self.dim_tags_set_implicit_only_wrapped)

  @property
  def dim_tags_set_implicit(self):
    """
    This is mostly intended to be used for verification, such as ``out_shape`` in a layer.
      https://github.com/rwth-i6/returnn/issues/706

    We return a set because when dim tags (dimensions, and the shape) are checked,
    we never want that the order plays any role.
      https://github.com/rwth-i6/returnn/wiki/RETURNN-principles
    Further, dimension tags should ideally be unique.
      https://github.com/rwth-i6/returnn/issues/632
    (This is not enforced currently, but we should not treat this specially now.)

    :return: set of dim tags
    :rtype: set[Dim]
    """
    dims = set(self.dim_tags)
    dims.update(self.dim_tags_set_implicit_only)
    return dims

  @property
  def ndim(self):
    """
    :rtype: int
    :return: ndim counted without batch-dim
    """
    return len(self.shape)

  @property
  def ndim_dense(self):
    """
    :rtype: int
    :return: ndim counted without batch-dim, added by 1 if we are sparse
    """
    if self.sparse:
      return self.ndim + 1
    return self.ndim

  @property
  def batch_ndim(self):
    """
    :rtype: int
    :return: ndim counted with batch-dim
    """
    return len(self._dim_tags)

  @property
  def batch_ndim_dense(self):
    """
    :rtype: int
    :return: ndim counted with batch-dim, added by 1 if we are sparse
    """
    if self.sparse:
      return self.batch_ndim + 1
    return self.batch_ndim

  @property
  def is_time_major(self):
    """
    :return: whether this is in time-major format, i.e. (time,batch,...)
    :rtype: bool
    """
    return self.time_dim_axis == 0

  @property
  def is_batch_major(self):
    """
    :return: whether this is in batch-major format, i.e. (batch,...)
    :rtype: bool
    """
    return self.batch_dim_axis == 0

  @property
  def is_batch_feature_major(self):
    """
    :return: whether this is in batch-feature-major format, i.e. (batch,feature,...) (NC...)
    :rtype: bool
    """
    return self.batch_dim_axis == 0 and self.feature_dim_axis == 1

  @property
  def batch_dim_axis(self):
    """
    :return: batch dim axis, counted with batch-dim
    :rtype: int|None
    """
    return _batch_dim_axis_from_dim_tags_tuple(self._dim_tags)

  @batch_dim_axis.setter
  def batch_dim_axis(self, axis):
    """
    :param int|None axis:
    """
    if axis == self.batch_dim_axis:
      return
    raise Exception("%s: cannot set batch_dim_axis = %s" % (self, axis))

  def _default_feature_dim_axis(self):
    """
    :return: feature dim axis, counted with batch-dim
    :rtype: int|None
    """
    return _default_feature_dim_axis(
      batch_dim_axis=self.batch_dim_axis, time_dim_axis=self.time_dim_axis,
      batch_shape=self.batch_shape, sparse=self.sparse)

  @property
  def feature_dim_axis(self):
    """
    :return: feature dim axis, counted with batch-dim
    :rtype: int|None
    """
    if self._feature_dim_axis is not NotSpecified:
      return self._feature_dim_axis
    return self._default_feature_dim_axis()

  @feature_dim_axis.setter
  def feature_dim_axis(self, value):
    """
    :param int|None|NotSpecified value:
    """
    assert value is NotSpecified or value is None or isinstance(value, int)
    if isinstance(value, int):
      assert 0 <= value < self.batch_ndim
    self._feature_dim_axis = value

  @property
  def feature_dim_axis_or_unspecified(self):
    """
    :return: feature dim axis, counted with batch-dim. could also be unspecified
    :rtype: int|None|NotSpecified
    """
    return self._feature_dim_axis

  @property
  def time_dim_axis_excluding_batch(self):
    """
    :rtype: int|None
    """
    if self.time_dim_axis is None:
      return None
    return self.get_batch_axis_excluding_batch(self.time_dim_axis)

  @property
  def placeholder(self):
    """
    :rtype: tf.Tensor|None
    """
    return self._placeholder

  @placeholder.setter
  def placeholder(self, value):
    """
    :param tf.Tensor|None value:
    """
    self._placeholder = value
    self.sanity_check(assume_complete=False)

  @property
  def batch(self):
    """
    :rtype: BatchInfo|None
    """
    return self._batch

  @batch.setter
  def batch(self, batch):
    """
    :param BatchInfo|None batch:
    """
    if batch:
      assert batch.beam == self.beam
    self._batch = batch
    self._adapt_batch_consistent_dim_tags()

  @property
  def beam(self):
    """
    :rtype: SearchBeam|None
    """
    if self._beam:
      return self._beam
    if self._batch:
      return self._batch.beam
    return None

  @beam.setter
  def beam(self, beam):
    """
    :param SearchBeam|None beam:
    """
    # No check for batch.beam, as the batch is usually set only later.
    self._beam = beam
    if self._batch:
      self._batch = self._batch.copy_set_beam(beam=beam)
      self._adapt_batch_consistent_dim_tags()

  @property
  def dim(self):
    """
    :rtype: int|None
    """
    tag = self.feature_dim_or_sparse_dim
    if tag:
      return tag.dimension
    return None

  @dim.setter
  def dim(self, dim):
    """
    It is deprecated to explicitly set this.
    We just have this here to support some legacy code.
    It does nothing but checks the validity.

    :param int|None dim:
    """
    assert dim == self.dim

  @property
  def feature_dim_or_sparse_dim(self):
    """
    :return: if we have a feature dim, return its dim tag. if we are sparse, return the sparse_dim. otherwise None
    :rtype: Dim|None
    """
    if self.have_feature_axis():
      return self.dim_tags[self.feature_dim_axis]
    if self.sparse_dim:
      return self.sparse_dim
    return None

  @property
  def sparse(self):
    """
    :rtype: bool
    :return: whether the values represent class indices. see ``sparse_dim``
    """
    return bool(self.sparse_dim)

  @sparse.setter
  def sparse(self, sparse):
    """
    It is deprecated to explicitly set this.
    We just have this here to support some legacy code.

    :param bool sparse:
    """
    if self.sparse == sparse:
      return
    if not sparse:
      self.sparse_dim = None
      return
    raise Exception("%s: setting sparse=True not supported anymore. set sparse_dim instead" % self)

  @property
  def vocab(self):
    """
    :rtype: returnn.datasets.util.vocabulary.Vocabulary|None
    """
    if self.sparse_dim:
      return self.sparse_dim.vocab
    if self.have_feature_axis():
      return self.dim_tags[self.feature_dim_axis].vocab
    return None

  @vocab.setter
  def vocab(self, vocab):
    """
    :param returnn.datasets.util.vocabulary.Vocabulary|None vocab:
    """
    raise Exception("%s: setting vocab not supported anymore. set sparse_dim instead" % self)

  def time_dimension(self):
    """
    :return: shape(placeholder)[time_dim_axis], int scalar
    :rtype: tf.Tensor
    """
    from .basic import reuse_name_scope_of_tensor
    assert self.time_dim_axis is not None
    if self.batch_shape[self.time_dim_axis] is not None:
      return self.batch_shape[self.time_dim_axis]
    with reuse_name_scope_of_tensor(self.placeholder):
      with tf.name_scope("time_dim"):
        return tf.shape(self.placeholder)[self.time_dim_axis]

  def get_dim(self, axis):
    """
    :param int axis: counted with batch-dim
    :return: shape[axis]
    :rtype: tf.Tensor|int
    """
    if self.batch_shape[axis] is not None:
      return self.batch_shape[axis]
    return tf.shape(self.placeholder)[axis]

  def get_placeholder_as_time_major(self):
    """
    :rtype: tf.Tensor
    """
    assert self.placeholder is not None
    return self.copy_as_time_major().placeholder

  def get_placeholder_as_batch_major(self):
    """
    :rtype: tf.Tensor
    """
    assert self.placeholder is not None
    return self.copy_as_batch_major().placeholder

  def get_placeholder_with_specific_batch_dim_axis(self, batch_dim_axis):
    """
    :param int batch_dim_axis:
    :rtype: tf.Tensor
    """
    from .basic import swapaxes
    assert self.placeholder is not None
    if self.batch_dim_axis == batch_dim_axis:
      return self.placeholder
    return swapaxes(self.placeholder, batch_dim_axis, self.batch_dim_axis)

  def get_placeholder_with_runtime_sanity_checks(self):
    """
    :return: identity(self.placeholder) with added checks
    :rtype: tf.Tensor
    """
    with tf.control_dependencies([self.get_runtime_sanity_check_op()]):
      return tf.identity(self.placeholder, name="identity_with_runtime_sanity_checks")

  def get_placeholder_time_flattened(self):
    """
    :return: via :func:`flatten_with_seq_len_mask`
    :rtype: tf.Tensor
    """
    from .basic import flatten_with_seq_len_mask
    assert self.placeholder is not None
    assert self.have_time_axis()
    # flatten_with_seq_len_mask only works if either time_dim_axis or batch_dim_axis is 0:
    assert 0 in [self.time_dim_axis, self.batch_dim_axis]
    seq_lens = self.size_placeholder[self.time_dim_axis_excluding_batch]
    return flatten_with_seq_len_mask(self.placeholder, seq_lens, batch_dim_axis=self.batch_dim_axis,
                                     time_dim_axis=self.time_dim_axis)

  def get_placeholder_flattened(self, keepdims=False):
    """
    :param bool keepdims: if set, it will add broadcast dimensions after the flattening behind the first axis
    :rtype: tf.Tensor
    :return: placeholder where all dynamic axes are flattened into a single axis.
      e.g. for the usual case (batch, time, dim), it becomes (batch'|time', dim),
      or (batch, time, height, dim) will also become (batch'|time', dim).
      with keep_dims, (batch, time, height, dim) will become (batch'|time', 1, 1, dim).
    """
    assert self.placeholder is not None
    x = self.placeholder
    orig_dyn_axes = self.get_spatial_batch_axes() + [self.batch_dim_axis]
    dyn_axes = list(orig_dyn_axes)
    if dyn_axes == [self.batch_dim_axis]:
      return x
    assert 0 in dyn_axes, "would need some transpose, not supported at the moment"
    assert len(dyn_axes) > 1
    orig_num_dyn_axes = len(dyn_axes)
    ndim = len(self.batch_shape)
    if self.have_time_axis():
      x = self.get_placeholder_time_flattened()
      removed_axis = max(self.time_dim_axis, self.batch_dim_axis)
      dyn_axes.remove(removed_axis)
      dyn_axes = [(i if (i < removed_axis) else (i - 1))
                  for i in dyn_axes]
      ndim -= 1
    if len(dyn_axes) > 1:
      shape = tf.shape(x)
      x = tf.reshape(
        x,
        [tf.reduce_prod([shape[i] for i in dyn_axes])] +
        [shape[i] for i in range(ndim) if i not in dyn_axes])
      dyn_axes = [0]
    assert dyn_axes == [0]
    if keepdims and orig_num_dyn_axes >= 2:
      for i in orig_dyn_axes:
        if i not in dyn_axes:
          x = tf.expand_dims(x, axis=i)
      x.set_shape([None] * self.batch_ndim)
    return x

  def get_axes(self, exclude_time=False, exclude_batch=False, exclude_feature=False):
    """
    :param bool exclude_time: will filter out the time-axis
    :param bool exclude_batch: will filter out the batch-axis
    :param bool exclude_feature: will filter out the feature-axis
    :return: list of axes, like `range(len(self.shape))`, calculated with batch dim.
    :rtype: list[int]
    """
    axes = list(range(len(self.batch_shape)))
    if exclude_time and self.time_dim_axis is not None:
      axes.pop(axes.index(self.time_dim_axis))
    if exclude_batch and self.batch_dim_axis is not None:
      axes.pop(axes.index(self.batch_dim_axis))
    if exclude_feature and self.feature_dim_axis is not None:
      axes.pop(axes.index(self.feature_dim_axis))
    return axes

  @classmethod
  def _verify_axis_int_from_description(cls, allow_int=NotSpecified):
    """
    Call this when you have the case that ``axis`` or ``axes``
    in :func:`get_axes_from_description` or :func:`get_axis_from_description`
    was specified as int.

    :param bool|NotSpecified allow_int:
    """
    msg = "Do not specify axis as int but as str or Dim instead."
    if allow_int is NotSpecified:
      from returnn.util import BehaviorVersion
      BehaviorVersion.require(condition=False, message=msg, version=5)
    if allow_int:
      return
    raise Exception(msg)

  @classmethod
  def _verify_axis_order_dependent(cls):
    """
    Call this when you have the case that ``axis`` or ``axes``
    in :func:`get_axes_from_description` or :func:`get_axis_from_description`
    depends on the order of the axes.
    """
    from returnn.util import BehaviorVersion
    BehaviorVersion.require(
      condition=False,
      message="Do not specify axis or axes in a way that depends on the order of the axes.",
      version=7)

  def _make_valid_int_axis(self, axis):
    """
    :param int axis: counted with batch. anything in [-ndim,ndim-1]
    :return: axis in [0,ndim-1]
    :rtype: int
    """
    if axis < 0:
      assert axis + self.batch_ndim >= 0
      axis += self.batch_ndim
    assert axis < self.batch_ndim
    return axis

  def get_axes_from_description(self, axes, allow_int=NotSpecified):
    """
    :param int|list[int]|str|typing.Sequence[str|Dim]|Dim|None axes: one axis or multiple axis, or none.
      This is counted with batch-dim, which by default is axis 0 (see enforce_batch_dim_axis).
      It also accepts the special tokens "B"|"batch", "spatial", "spatial_except_time", or "F"|"feature",
      and more (see the code).
    :param bool|NotSpecified allow_int: whether to allow an int directly.
      in almost all cases, it is better to use a symbolic name
      to specify an axis, as different layers could reorder them, and maybe also change their behavior in the future.
    :return: list of axes, counted with batch-dim
    :rtype: list[int]
    """
    if axes is None or axes == "":
      return []
    if isinstance(axes, Dim):
      # Once we have not guaranteed unique dim tags, multiple axes could match.
      # https://github.com/rwth-i6/returnn/issues/632
      dims = [i for (i, tag) in enumerate(self.dim_tags) if tag == axes]
      if len(dims) > 1:
        max_match_priority = max(self.dim_tags[i].match_priority for i in dims)
        dims = [i for i in dims if self.dim_tags[i].match_priority == max_match_priority]
      assert len(dims) <= 1, (
        "%s: matching dim %s must be unique,"
        " use `match_priority` to resolve the matching order of ambiguous dimensions" % (self, axes))
      return dims
    if isinstance(axes, int):
      self._verify_axis_int_from_description(allow_int=allow_int)
      return [self._make_valid_int_axis(axes)]
    assert isinstance(axes, (str, int, list, tuple, typing.Sequence))
    if isinstance(axes, str):
      import re
      axes = axes.lower()
      if axes in ["b", "batch"]:
        assert self.batch_dim_axis is not None
        return [self.batch_dim_axis]
      elif axes == "spatial":
        return self.get_spatial_batch_axes()
      elif re.match("(s|spatial):-?\\d+$", axes):
        self._verify_axis_order_dependent()
        s = int(axes.split(":")[1])
        spatial_axes = self.get_spatial_batch_axes()
        if s < 0:
          s += len(spatial_axes)
        assert s < len(spatial_axes), "%s get_axes_from_description: %r invalid" % (self, axes)
        return [spatial_axes[s]]
      elif axes in ["dyn", "dynamic"]:
        return self.get_dynamic_axes()
      elif re.match("(d|dyn|dynamic):-?\\d+$", axes):
        self._verify_axis_order_dependent()
        s = int(axes.split(":")[1])
        dyn_axes = self.get_dynamic_axes()
        if s < 0:
          s += len(dyn_axes)
        assert 0 <= s < len(dyn_axes), "%s get_axes_from_description: %r invalid" % (self, axes)
        return [dyn_axes[s]]
      elif axes == "spatial_except_time":
        axes = self.get_spatial_batch_axes()
        assert self.time_dim_axis is not None
        axes.remove(self.time_dim_axis)
        return axes
      elif axes in ["t", "time"]:
        assert self.time_dim_axis is not None
        return [self.time_dim_axis]
      elif axes == "t?":
        return [self.time_dim_axis] if self.time_dim_axis is not None else []
      elif axes == "except_time":  # also except batch
        axes = list(range(self.batch_ndim))
        axes.remove(self.batch_dim_axis)
        if self.time_dim_axis is not None:
          axes.remove(self.time_dim_axis)
        return axes
      elif axes == "except_batch":
        axes = list(range(self.batch_ndim))
        axes.remove(self.batch_dim_axis)
        return axes
      elif re.match("(except_batch):-?\\d+$", axes):
        self._verify_axis_order_dependent()
        s = int(axes.split(":")[1])
        non_batch_axes = list(range(self.batch_ndim))
        if self.batch_dim_axis is not None:
          non_batch_axes.remove(self.batch_dim_axis)
        if s < 0:
          s += len(non_batch_axes)
        assert 0 <= s < len(non_batch_axes), "%s get_axes_from_description: %r invalid" % (self, axes)
        return [non_batch_axes[s]]
      elif axes == "*":
        return list(range(self.batch_ndim))
      elif axes == "static":
        return self.get_static_axes()
      elif re.match("(static):-?\\d+$", axes):
        self._verify_axis_order_dependent()
        s = int(axes.split(":")[1])
        static_axes = self.get_static_axes()
        if s < 0:
          s += len(static_axes)
        assert 0 <= s < len(static_axes), "%s get_axes_from_description: %r invalid" % (self, axes)
        return [static_axes[s]]
      elif re.match("(dim):\\d+$", axes):
        s = int(axes.split(":")[1])
        dims = [a for a in range(self.batch_ndim) if self.batch_shape[a] == s]
        assert len(dims) == 1, "%s get_axes_from_description: dim %i only allowed when unique" % (self, s)
        return dims
      elif axes in ["f", "feature", "non_spatial"]:
        return self.get_feature_batch_axes()
      elif all([a in "btf" for a in axes]):
        return self.get_axes_from_description(list(axes))
      elif axes.startswith("stag:"):  # spatial tag
        return [self.get_axis_by_tag_name(axes[len("stag:"):], spatial_only=True)]
      elif axes.startswith("stag-single:"):  # spatial tag which possibly matches multiple spatial axes
        # in this case, a name of form "stag-single:<idx>:<name> is expected.
        # idx is relative to the matching stags, i.e., it is the index among the list of spatial dims matching the name
        # Note: no _verify_axis_order_dependent here because as long as we do not enforce unique dim tags
        # (https://github.com/rwth-i6/returnn/issues/632), we can have multiple axes with the same tag,
        # and then we need to be able to differentiate between them by order.
        _, idx_s, name = axes.split(":", 2)  # stag-single:<idx>:<name>
        idx = int(idx_s)
        return [self.get_axes_by_tag_name(name, spatial_only=True)[idx]]
      raise Exception("invalid axis mode %r" % axes)
    assert isinstance(axes, (tuple, list, typing.Sequence)), "invalid axes %r" % axes
    flat_axes = []
    for i in axes:
      if isinstance(i, int):
        self._verify_axis_int_from_description(allow_int=allow_int)
        flat_axes.append(self._make_valid_int_axis(i))
      else:
        assert isinstance(i, (str, tuple, list, Dim))
        flat_axes += self.get_axes_from_description(i, allow_int=allow_int)
    res = []
    for i in flat_axes:
      if i not in res:
        res.append(i)
    return res

  def get_dim_tag_from_description(self, axis):
    """
    :param str|Dim axis:
    :return: our matching dim tag. this assumes it exists.
    :rtype: Dim
    """
    axis_int = self.get_axis_from_description(axis, allow_int=False)
    return self.dim_tags[axis_int]

  def get_axis_from_description(self, axis, allow_int=NotSpecified):
    """
    :param int|str|Dim axis:
    :param bool|NotSpecified allow_int:
    :return: axis, counted with batch-dim
    :rtype: int
    """
    axes = self.get_axes_from_description(axis, allow_int=allow_int)
    assert axes, "%s: %r axis not found" % (self, axis)
    assert len(axes) == 1, "%r: %r is not a unique axis but %r" % (self, axis, axes)
    return axes[0]

  def get_description_from_axis(self, axis):
    """
    :param int axis:
    :return: some canonical description, such that ``self.get_axis_from_description(res) == axis``.
      This is quite heuristically for now. We use both strings as also Dim when appropriate.
      The behavior could potentially change in the future, also the condition will always hold.
    :rtype: str|Dim
    """
    assert 0 <= axis < self.batch_ndim
    if axis == self.batch_dim_axis:
      return "B"
    dim_tag = self.dim_tags[axis]
    # It's possible that dim tags are not unique (https://github.com/rwth-i6/returnn/issues/632).
    matching_tags = [i for (i, tag) in enumerate(self.dim_tags) if tag == dim_tag]
    if dim_tag.dyn_size_ext and len(matching_tags) == 1:
      return dim_tag
    if axis == self.time_dim_axis:
      return "T"  # this might change
    if axis == self.feature_dim_axis:
      return "F"  # this might change
    if len(matching_tags) == 1:
      # Fallback with dim tag
      return dim_tag
    # Do not use indexed static or dynamic because we want to avoid relying on the axis order as much as possible.
    # However, as we do not have unique dim tags in this case, we have to rely at least on the order of this dim tag.
    # Use stag-single.
    name = dim_tag.description
    matching_axes = self.get_axes_by_tag_name(name, spatial_only=True)
    assert axis in matching_axes
    return (
      "stag-single:%i:%s" % (
        matching_axes.index(axis) - len(matching_axes), name))  # negative because this is likely more robust

  def has_axis(self, axis):
    """
    :param str|Dim axis:
    :return: whether the axis exists
    :rtype: bool
    """
    axes = self.get_axes_from_description(axis, allow_int=False)
    return len(axes) > 0

  def get_axes_by_tag_name(self, name, spatial_only=False):
    """
    :param str name: the tag name, or part of it (must be unique, and must exist)
    :param bool spatial_only:
    :rtype: list[int]
    """
    dim_tags = self.get_batch_shape_dim_tags()
    matching_dim_tags = [
      (axis, tag) for axis, tag in enumerate(dim_tags)
      if name.lower() in tag.description.lower()
      or name.lower() in tag.get_same_base().description.lower()]
    if spatial_only:
      spatial_axes = self.get_spatial_batch_axes()
      matching_dim_tags = [
        (axis, tag) for axis, tag in matching_dim_tags if axis in spatial_axes or tag.is_spatial_dim()]
    return [ax for ax, _ in matching_dim_tags]

  def get_axis_by_tag_name(self, name, spatial_only=False):
    """
    :param str name: the tag name, or part of it (must be unique, and must exist)
    :param bool spatial_only:
    :rtype: int
    """
    matching_dim_tags = self.get_axes_by_tag_name(name, spatial_only)
    assert len(matching_dim_tags) > 0, "%r: no %stag found with name %r" % (
      self, "spatial " if spatial_only else "", name)
    assert len(matching_dim_tags) == 1, "%r: tag name %r is not unique in dim tags %r" % (
      self, name, self.get_batch_shape_dim_tags())
    return matching_dim_tags[0]

  def get_batch_axis_excluding_batch(self, axis):
    """
    :param int axis: counted with batch-dim
    :return: axis counted without batch-dim
    :rtype: int|None
    """
    return _get_axis_wo_b(axis, batch_dim_axis=self.batch_dim_axis, batch_ndim=self.batch_ndim)

  def have_dim_tag(self, tag, include_implicit=True, unique=False):
    """
    :param Dim tag:
    :param bool include_implicit:
    :param bool unique:
    :rtype: bool
    """
    dims = list(self.dim_tags)
    if include_implicit:
      dims.extend(self.dim_tags_set_implicit_only)
    matching_dims = [dim for dim in dims if dim == tag]
    return (len(matching_dims) == 1) if unique else (len(matching_dims) >= 1)

  def get_batch_axis(self, axis):
    """
    :param int axis: counted without batch-dim
    :return: axis counted with batch-dim
    :rtype: int
    """
    return _get_axis_wb(axis, batch_dim_axis=self.batch_dim_axis)

  def have_batch_axis(self):
    """
    :rtype: bool
    """
    return self.batch_dim_axis is not None

  def have_time_axis(self):
    """
    :rtype: bool
    """
    return self.time_dim_axis is not None

  def have_feature_axis(self):
    """
    :rtype: bool
    """
    return self.feature_dim_axis is not None

  def is_time_axis_dynamic(self):
    """
    :return: whether there are different seq-lens for the time, or all the same (static)
    :rtype: bool
    """
    assert self.time_dim_axis is not None
    if self.placeholder is None:
      # Run at template construction time.
      return self.batch_shape[self.time_dim_axis_excluding_batch] is None
    if self.time_dim_axis_excluding_batch in self.size_placeholder:
      return True
    assert isinstance(self.shape[self.time_dim_axis_excluding_batch], int), (
      "%s: dynamic time axis dim (None) (axis %i) but size_placeholder %r misses information" % (
        self, self.time_dim_axis, self.size_placeholder))
    return False

  def is_axis_dynamic(self, axis):
    """
    :param int axis: counted with batch-dim axis
    :return: dynamic, i.e. we have it in size_placeholder.
      Note that this does not perfectly match with :func:`get_dynamic_axes`, but more with :func:`is_time_axis_dynamic`,
      although probably in most (all?) cases it should match.
      If True, you can get the size via :func:`get_dynamic_size`.
    :rtype: bool
    """
    if axis == self.batch_dim_axis:
      return False
    return self.batch_shape[axis] is None

  def has_dynamic_size(self, axis):
    """
    :param int axis: counted with batch-dim axis. implies that you can call :func:`get_dynamic_size`.
    :rtype: bool
    """
    return self.dim_tags[axis].dyn_size is not None

  def get_dynamic_size(self, axis):
    """
    :param int axis: counted with batch-dim axis. :func:`get_dynamic_size` should be True.
    :return: shape (B,)
    :rtype: tf.Tensor
    """
    tag = self.dim_tags[axis]
    assert tag.dyn_size is not None, "%s: axis %i has no dyn size" % (self, axis)
    return tag.dyn_size

  def set_dynamic_size(self, axis, sizes):
    """
    :param int axis: counted with batch-dim
    :param tf.Tensor sizes: shape [B]
    """
    # Note: The following code is somewhat ugly patchwork
    # to fix some other currently incomplete or buggy behavior of some layers
    # which introduce sizes without correctly setting the dim tag.
    # The beam information is also missing currently.
    # We make the ugly assumption that when it is unset,
    # the first usage should hopefully define the correct beam.
    if getattr(sizes, "_RETURNN_dyn_size_beam", NotSpecified) is NotSpecified:
      sizes._RETURNN_dyn_size_beam = self.beam
    if self.beam and getattr(sizes, "_RETURNN_dyn_size_beam", None) != self.beam:
      tag = Dim.get_tag_from_size_tensor(sizes)
      assert tag and self.batch
      tag = tag.get_for_batch_ctx(batch=self.batch, ctx=self.control_flow_ctx)
      assert tag.dyn_size is not None
      sizes = tag.dyn_size

    sizes_tag = Dim.get_tag_from_size_tensor(sizes)
    if sizes_tag:
      assert sizes_tag.is_same_size_tensor(sizes)
    tag = self.dim_tags[axis]
    assert tag.dimension is None  # dynamic axis
    if tag.is_same_size_tensor(sizes):
      return  # nothing to do
    if tag.dyn_size is None:
      if sizes_tag:  # special rule for older code: overtake previous existing
        assert sizes_tag.is_same_size_tensor(sizes)
        self._dim_tags = self.dim_tags[:axis] + (sizes_tag,) + self.dim_tags[axis + 1:]
        # Also assume the existing dim tag should be expected as equal.
        # Likely there is anyway no reference so this does not matter.
        tag.declare_same_as(sizes_tag)
      else:
        # Assign now. This should also set the dim tag on sizes.
        new_tag = tag.set_tag_on_size_tensor(sizes, batch=self.batch)
        if new_tag is not tag:
          self._dim_tags = self.dim_tags[:axis] + (new_tag,) + self.dim_tags[axis + 1:]
    else:
      # Reset to some new size.
      # Use new dim tag, or previous existing attached to size.
      assert sizes_tag, "%s: assign dyn sizes %s without defined dim tag" % (self, sizes)
      self._dim_tags = self.dim_tags[:axis] + (sizes_tag,) + self.dim_tags[axis + 1:]

  def get_dynamic_axes(self):
    """
    :return: list of axes, counted with batch-dim axis (but we exclude the batch dim axis itself)
    :rtype: list[int]
    """
    return [axis for axis, dim in enumerate(self.batch_shape)
            if axis != self.batch_dim_axis and dim is None]

  def get_static_axes(self):
    """
    :return: list of axes, counted with batch-dim axis (but we exclude the batch dim axis itself)
    :rtype: list[int]
    """
    return [axis for axis, dim in enumerate(self.batch_shape)
            if axis != self.batch_dim_axis and dim is not None]

  def mark_same_time(self, tags, must_match=False):
    """
    If the given dimension tag matches any of our axes, we set our time axis to the selected one.

    :param set[Dim]|Dim tags:
    :param bool must_match: if True, throw an exception if not found
    :return: whether we have found the same
    :rtype: bool
    """
    if isinstance(tags, Dim):
      tags = {tags}
    assert all(isinstance(tag, Dim) for tag in tags)
    for axis, dim_tag in enumerate(self.dim_tags):
      if dim_tag in tags:
        self.time_dim_axis = axis
        return True
    if must_match:
      raise Exception("%s mark_same_time: %s not found" % (self, tags))
    return False

  def is_same_time_dim(self, other):
    """
    Checks whether we have a matching/compatible time dim.

    :param Data other:
    :rtype: bool
    """
    assert self.have_time_axis()
    if not other.have_time_axis():
      return False
    tag_self = self.get_dim_tag(self.time_dim_axis)
    tag_other = other.get_dim_tag(other.time_dim_axis)
    return tag_self == tag_other

  def get_sequence_lengths(self):
    """
    :return: seq lens tensor of shape (batch,) of dtype int32. also see :func:`get_dynamic_size`
    :rtype: tf.Tensor
    """
    from .basic import same_control_flow_ctx, expand_dims_unbroadcast
    assert self.time_dim_axis is not None
    if self.is_time_axis_dynamic():
      return self.size_placeholder[self.time_dim_axis_excluding_batch]
    assert self.shape[self.time_dim_axis_excluding_batch] is not None
    with same_control_flow_ctx(self.placeholder), tf.name_scope("fixed_seq_len"):
      return expand_dims_unbroadcast(
        self.shape[self.time_dim_axis_excluding_batch], axis=0, dim=self.get_batch_dim())

  def get_sequence_mask(self):
    """
    :return: seq mask of shape (batch,time) if we are batch-major, else (time,batch) if we are time-major
    :rtype: tf.Tensor
    """
    from .basic import sequence_mask_time_major, sequence_mask
    assert self.time_dim_axis is not None
    assert self.batch_dim_axis is not None
    if self.is_time_major:
      assert self.batch_dim_axis == 1
      return sequence_mask_time_major(self.get_sequence_lengths())
    else:
      assert self.batch_dim_axis == 0
      assert self.time_dim_axis == 1
      return sequence_mask(self.get_sequence_lengths())

  def get_sequence_mask_broadcast(self, axis=None):
    """
    :param int|None axis:
    :return: seq mask of shape ((batch,time) or (time,batch)) + (1,)s for remaining dims
      if BT or TB major, and axis is T or None.
      In general compatible to placeholder, i.e. same ndim, with broadcast dims.
      We assert here that the axis is dynamic (:func:`is_axis_dynamic`), i.e. we have the size.
    :rtype: tf.Tensor
    """
    from .basic import sequence_mask_time_major, sequence_mask
    if axis is None:
      assert self.time_dim_axis is not None
      axis = self.time_dim_axis
    if axis < 0:
      assert axis + self.batch_ndim > 0
      axis += self.batch_ndim
    assert 0 <= axis < self.batch_ndim
    assert axis != self.batch_dim_axis
    tag = self.dim_tags[axis]
    assert tag.dyn_size_ext
    assert set(tag.dyn_size_ext.dim_tags).issubset(self.dim_tags)  # https://github.com/rwth-i6/returnn/issues/721
    with tf.name_scope("get_sequence_mask_broadcast"):
      if tag.dyn_size_ext.have_batch_axis() and tag.dyn_size_ext.batch_ndim == 1:  # just [B]
        # This is the common case where the size is of shape [B].
        # We make use of sequence_mask or sequence_mask_time_major in that case,
        # which is optimized by caching.
        size = tag.dyn_size
        if axis >= self.batch_dim_axis:
          seq_mask = sequence_mask(size)  # (B,T)
        else:  # axis < batch_dim_axis
          seq_mask = sequence_mask_time_major(size)  # (T,B)
        shape = [1] * self.batch_ndim  # type: typing.List[typing.Union[int,tf.Tensor]]
        placeholder_shape = tf.shape(self.placeholder)
        shape[self.batch_dim_axis] = placeholder_shape[self.batch_dim_axis]
        shape[axis] = placeholder_shape[axis]
        seq_mask = tf.reshape(seq_mask, shape, name="seq_mask_reshape")
        assert seq_mask.get_shape().ndims == self.batch_ndim
      else:  # size is something unusual, not just [B], but e.g. [B,S] or so
        max_idx = tf.reduce_max(tag.dyn_size)
        # We use the assumption that self.placeholder.shape[axis] == max_idx.
        idx_range = tf.range(max_idx)
        idx_range = tf.reshape(idx_range, [1] * axis + [max_idx] + [1] * (self.batch_ndim - axis - 1))
        # size_ext might have invalid (zero) sizes when it itself has some padding, e.g. when its own shape is dynamic.
        # A zero size can lead to problems in some cases, e.g. in SoftmaxOverSpatialLayer,
        # when everything is masked to -inf, it results in nan, and this likely produces nan in backprop or elsewhere.
        # Thus, mask size_ext itself, and set the padded values to 1.
        # This assumes that max_idx >= 1.
        size_ext = tag.dyn_size_ext.copy_masked(max_idx)
        size_ext = size_ext.copy_compatible_to(self, check_sparse=False, check_dtype=False)
        seq_mask = tf.less(idx_range, size_ext.placeholder)
        assert seq_mask.get_shape().ndims == self.batch_ndim
    return seq_mask

  def get_sequence_lengths_broadcast(self, axis=None):
    """
    :param int|None axis:
    :return: seq len of some shape which is broadcastable to self.placeholder.
      Note that this is not always possible, e.g. when the seq len has shape [B]
      but the tensor has just shape [T]. We currently throw an error then.
    :rtype: tf.Tensor
    """
    if axis is None:
      assert self.time_dim_axis is not None
      axis = self.time_dim_axis
    if axis < 0:
      assert axis + self.batch_ndim > 0
      axis += self.batch_ndim
    assert 0 <= axis < self.batch_ndim
    assert axis != self.batch_dim_axis
    tag = self.dim_tags[axis]
    assert tag.dyn_size_ext
    return tag.dyn_size_ext.copy_compatible_to(self, check_dtype=False, check_sparse=False).placeholder

  def copy_masked(self, mask_value):
    """
    :param float|int|tf.Tensor mask_value:
    :rtype: Data
    """
    assert self.placeholder is not None
    from .basic import mask_dyn_seq_len_nd
    dyn_axes = [axis for axis, dim in enumerate(self.dim_tags) if not dim.is_batch_dim() and dim.dimension is None]
    res = self.copy()
    res.placeholder = mask_dyn_seq_len_nd(self, pad_value=mask_value, axes=dyn_axes)
    return res

  def get_batch_dim(self):
    """
    :rtype: tf.Tensor|int
    """
    assert self.batch_dim_axis is not None
    if self.batch:
      if self.beam:
        assert self.batch.beam == self.beam
      dim = self.batch.dim
      if isinstance(dim, tf.Tensor):
        batch_dim_ = self.dim_tags[self.batch_dim_axis]
        batch_dim_.set_tag_on_size_tensor(dim, batch=self.batch)
      return dim
    # Note: We need this fallback code for now
    # until we consistently have set self.batch correctly in all cases.
    from returnn.tf.layers.base import LayerBase
    batch = LayerBase.get_recent_layer().get_batch_info()
    batch = batch.copy_set_beam(self.beam)
    return batch.dim

  def get_batch_dim_tag(self):
    """
    :rtype: Dim
    """
    assert self.have_batch_axis()
    return self.dim_tags[self.batch_dim_axis]

  def get_static_batch_dim(self):
    """
    :rtype: int|None
    """
    # Do not fallback to get_batch_dim or get_recent_layer or so. This should be safe.
    if not self.batch:
      return None
    return self.batch.static_dim

  def get_spatial_batch_axes(self):
    """
    :rtype: list[int]
    :return: list of axes which are not batch axes and not feature or which are time axis or dynamic.
      counted with batch-dim.
    """
    return [
      axis
      for axis in range(self.batch_ndim)
      if axis != self.batch_dim_axis
      and (axis != self.feature_dim_axis or
           axis == self.time_dim_axis or
           self.batch_shape[axis] is None)]

  def get_spatial_axes(self):
    """
    :rtype: list[int]
    :return: list of axes which are not feature and batch axes, counted without batch-dim.
    """
    return [self.get_batch_axis_excluding_batch(axis) for axis in self.get_spatial_batch_axes()]

  def get_feature_batch_axes(self):
    """
    :rtype: list[int]
    :return: list of axes which are feature axes, counted with batch-dim. currently there is only one or zero such axis.
    """
    if self.feature_dim_axis is not None:
      return [self.feature_dim_axis]
    return []

  def get_feature_axes(self):
    """
    :rtype: list[int]
    :return: list of axes which are feature axes, counted without batch-dim.
    """
    return [self.get_batch_axis_excluding_batch(axis) for axis in self.get_feature_batch_axes()]

  # Exclude "batch_dim_axis" now because that is always inferred from dim tags.
  SpecialAxesNames = ("time_dim_axis", "feature_dim_axis")

  def get_special_axes_dict(self, counted_with_batch_dim=True, only_available=False):
    """
    :param bool counted_with_batch_dim:
    :param bool only_available:
    :return: dict axis-name -> axis
    :rtype: dict[str,int]
    """
    axes = list(self.SpecialAxesNames)
    d = {k: getattr(self, k) for k in axes}
    if not counted_with_batch_dim:
      d = {k: self.get_batch_axis_excluding_batch(v) if (v is not None) else None
           for (k, v) in d.items()}
    if only_available:
      d = {k: v for (k, v) in d.items() if v is not None}
      if self._feature_dim_axis is NotSpecified:  # special rule
        d.pop("feature_dim_axis", None)
    return d

  def get_bc_spatial_batch_shape(self):
    """
    :return: shape which will broadcast along all spatial dimensions and time/batch dim
    :rtype: tuple[int|None]
    """
    dyn_axes = self.get_spatial_batch_axes()
    if self.batch_dim_axis is not None:
      dyn_axes += [self.batch_dim_axis]
    return tuple([1 if (axis in dyn_axes) else dim
                  for axis, dim in enumerate(self.batch_shape)])

  def get_bc_shape(self, opts=None):
    """
    :param dict[Dim|str|list[Dim|str]|tuple[Dim|str],int|str|None]|None opts:
      ``key`` specifies the axes.
      ``value`` 1 ('x') is broadcasting, -1 (None) is not broadcasting
      Axes should not be defined multiple times.
      The default behavior if an axis is not specified is like :func:`get_bc_spatial_batch_shape`,
      i.e. it will broadcast in batch and spatial dims only.
      Or if "*" is in the dict, this overwrites the default behavior for all axes.
    :return: shape where 1 means broadcasting, None or >1 means not broadcasting. can be used for :func:`TFUtil.dropout`
    :rtype: tuple[int|None]
    """
    if opts is None:
      opts = {}
    default_axes_map = dict(enumerate(self.get_bc_spatial_batch_shape()))
    axes_map = {}  # int -> int|None
    for key, value in opts.items():
      assert value in (-1, 1, 'x', None), "%r get_bc_shape: invalid value in opts %r" % (self, opts)
      if value == 'x':
        value = 1
      if value == -1:
        value = None
      key_axes = self.get_axes_from_description(key)
      for key_axis in key_axes:
        assert key_axis not in axes_map, (
          "%r get_bc_shape: axis %i is defined multiple times in opts %r" % (self, key_axis, opts))
        assert 0 <= key_axis < self.batch_ndim, "%r get_bc_shape: invalid axis %i in opts %r" % (self, key_axis, opts)
        (axes_map if key != "*" else default_axes_map)[key_axis] = (
          self.batch_shape[key_axis] if value is None else value)
    # Fill in remaining axes by defaults, just as in get_bc_spatial_batch_shape.
    remaining_axes = sorted(set(range(self.batch_ndim)).difference(axes_map.keys()))
    for axis in remaining_axes:
      axes_map[axis] = default_axes_map[axis]
    assert sorted(axes_map.keys()) == list(range(self.batch_ndim))
    return tuple([axes_map[i] for i in range(self.batch_ndim)])

  def get_scope_name(self):
    """
    :return: via self.placeholder or any self.size_placeholder, or None
    :rtype: str|None
    """
    if self.placeholder is not None:
      return os.path.dirname(self.placeholder.name)
    if self.size_placeholder:
      for i, v in sorted(self.size_placeholder.items()):
        if v is not None:
          return os.path.dirname(v.name)
    return None

  def get_full_name(self):
    """
    :return: if we have a defined scope (via :func:`self.get_scope_name`), then scope_name + "/" + self.name,
      otherwise just self.name
    :rtype: str
    """
    scope_name = self.get_scope_name()
    if scope_name:
      return "%s/%s" % (scope_name, self.name)
    return self.name

  def get_dim_tag(self, axis):
    """
    :param int axis: counted with batch-dim
    :rtype: Dim
    """
    return self._dim_tags[axis]

  def get_time_dim_tag(self):
    """
    :rtype: Dim
    """
    assert self.time_dim_axis is not None
    return self.get_dim_tag(self.time_dim_axis)

  def get_size_dim_tag(self, number):
    """
    :param int number: index in sorted(size_placeholder.keys())
    :rtype: Dim
    """
    axis_wo_batch = sorted(self.size_placeholder.keys())[number]
    return self.get_dim_tag(self.get_batch_axis(axis_wo_batch))

  def get_batch_shape_dim_tags(self):
    """
    :return: list of dimension tags, for each axis (counted with batch dim, i.e. len is batch_ndim)
    :rtype: tuple[Dim]
    """
    return self.dim_tags

  @classmethod
  def get_common_data(cls, sources, ignore_feature_dim=False, allow_broadcast_all_sources=NotSpecified, name=None):
    """
    :param list[Data] sources:
    :param bool ignore_feature_dim: when set, the feature dim does not have to match in the sources
    :param bool|NotSpecified allow_broadcast_all_sources:
    :param str|None name:
    :return: some generic data where the sources should be compatible to (with copy_compatible_to),
      i.e. it contains the union of all axes from all sources (least common multiple).
      This is always a template, and a new copy.
    :rtype: Data|None
    """
    from returnn.util import BehaviorVersion
    if not sources:
      return None
    assert sources
    if len(sources) == 1:
      return sources[0].copy_template()
    max_ndim = max([s.batch_ndim for s in sources])
    common_batch = BatchInfo.get_common_batch_info([src.batch for src in sources if src.batch])
    # Try with the (first) largest.
    common = [s for s in sources if s.batch_ndim == max_ndim][0]
    common = common.copy_template(name=name)
    common.beam = None  # this will be reset
    if common_batch:
      common.batch = common_batch.copy_set_beam(None)  # the beam will be reset
    if any([s.beam for s in sources]):
      # Note: we don't use copy_extend_with_beam because we don't want to create any ops in the TF graph at this point.
      common.beam = SearchBeam.get_combined_beam(*[s.beam for s in sources])
    is_equal_opts = dict(
      ignore_feature_dim=ignore_feature_dim, treat_feature_as_spatial=True,
      allow_same_spatial_dim=True,
      undefined_matches=True, derived_matches=True)
    if BehaviorVersion.get() < 11:
      is_equal_opts["broadcast_matches"] = True
    all_dim_tags, tags_dict = Dim.get_all_dimension_tags(sources, is_equal_opts=is_equal_opts)
    # Check for missing tags, and add those.
    for dim_tag in all_dim_tags:
      common_tag = Dim.get_existing_tag_from_collection(dim_tag, common.dim_tags, is_equal_opts=is_equal_opts)
      if common_tag:
        # Already have this tag. However, maybe we have a better one.
        # Dim.get_all_dimension_tags() would have selected that.
        if dim_tag != common_tag:
          axis = common.dim_tags.index(common_tag)
          common = common.copy_template_replace_dim_tag(axis=axis, new_dim_tag=dim_tag)
      else:
        axis = common.get_default_new_axis_for_dim_tag(dim_tag)
        common = common.copy_add_dim_by_tag(dim_tag, unbroadcast=True, axis=axis)
    if all(s.batch_ndim < common.batch_ndim for s in sources):
      from .basic import validate_broadcast_all_sources
      validate_broadcast_all_sources(
        allow_broadcast_all_sources=allow_broadcast_all_sources, inputs=sources, common=common)
    return common

  def find_matching_dims(self, dim_tag, is_equal_opts):
    """
    Finds the dimensions of this Data that match another Dim

    :param Dim dim_tag:
    :param dict[str,bool]|None is_equal_opts: passed to Dim.is_equal
    :rtype: list[int]
    :return: a list of matching axes, counted with batch dim. Sorted in ascending order
    """
    return [axis for axis in range(self.batch_ndim) if self.get_dim_tag(axis).is_equal(dim_tag, **is_equal_opts)]

  def find_matching_dim_map(self, other, other_axes, is_equal_opts=None):
    """
    Looks up all other_axes of another Data in this Data. Does not allow duplicates.

    :param Data other:
    :param list[int] other_axes: a list of axes of ``other``, counted with batch dim
    :return: a dict mapping other axes to own axes, all counted with batch dim
    :param dict[str,bool]|None is_equal_opts: passed to Dim.is_equal
    :rtype: dict[int,int]
    """
    if is_equal_opts is None:
      is_equal_opts = dict(
        allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True)

    def map_other_axis_to_self(other_axis, taken_self_axes):
      """
      :param int other_axis: counted with batch dim
      :param set[int] taken_self_axes: axes that should not be used again
      :return: the axis of ``self`` that matches ``other_axis``, counted with batch dim
      :rtype: int
      """
      other_axis_dim_tag = other.get_dim_tag(other_axis)
      is_equal_opts_ = None
      matching = None
      # First, try without any is_equal_opts. This is the most restrictive case.
      # Try with the given is_equal_opts.
      # Try harder by allowing broadcasting to match.
      # If still not, then also allow one single dyn_size to be unknown.
      for opt in [{}, is_equal_opts, "broadcast_matches", "unknown_spatial_matches"]:
        if isinstance(opt, dict):
          is_equal_opts_ = opt.copy()
        elif isinstance(opt, str):
          if opt in is_equal_opts_:
            continue
          is_equal_opts_[opt] = True
        matching = [
          self_axis for self_axis in self.find_matching_dims(other_axis_dim_tag, is_equal_opts_)
          if self_axis not in taken_self_axes]
        if opt == "unknown_spatial_matches":
          assert len(matching) <= 1, 'cannot match axes %s from %s to %s, failed at other %s, not unique after %s' % (
            other_axes, other, self, other_axis, opt)
        if matching:
          break
      assert matching, 'cannot match the axes %s from %s to %s. Failing at axis %s' % (
        other_axes, other, self, other_axis)
      # If there are multiple matches (e.g. because two axes have the same feature dim), leave their order intact.
      # We do this by always choosing the first unused match which is the smallest axes
      return matching[0]

    other_to_self_mapping = {}
    for axis in other_axes:
      other_to_self_mapping[axis] = map_other_axis_to_self(axis, set(other_to_self_mapping.values()))
    assert len(other_to_self_mapping) == len(other_axes), 'other_axes may not contain duplicates'
    return other_to_self_mapping


class _SizePlaceholderProxy:
  """
  This is a proxy object to emulate the original Data.size_placeholder behavior,
  which was a dict[int,tf.Tensor], axis_wo_batch -> sizes.
  """

  def __init__(self, data):
    """
    :param Data data:
    """
    self.data = data

  def _assert_sane_axis_wo_batch(self, idx):
    assert isinstance(idx, int) and 0 <= idx < self.data.ndim

  def __contains__(self, item):
    if not isinstance(item, int):
      return False
    if not 0 <= item < self.data.ndim:
      return False
    return self.data.has_dynamic_size(axis=self.data.get_batch_axis(item))

  def __getitem__(self, item):
    self._assert_sane_axis_wo_batch(item)
    return self.data.get_dynamic_size(axis=self.data.get_batch_axis(item))

  def __setitem__(self, key, value):
    self._assert_sane_axis_wo_batch(key)
    self.data.set_dynamic_size(axis=self.data.get_batch_axis(key), sizes=value)

  def __delitem__(self, key):
    self._assert_sane_axis_wo_batch(key)
    raise Exception("%s: cannot delete items from size_placeholder" % self.data)

  def __iter__(self):
    return iter(self.keys())

  def __len__(self):
    return len(self.keys())

  def __bool__(self):
    return bool(self.keys())

  __nonzero__ = __bool__  # Python 3 wants __bool__, Python 2.7 wants __nonzero__

  def __repr__(self):
    return repr(self.as_dict())

  def get(self, axis_wo_b, default=None):
    """
    :param int axis_wo_b:
    :param tf.Tensor|None default:
    :rtype: tf.Tensor|None
    """
    if axis_wo_b in self:
      return self[axis_wo_b]
    return default

  def pop(self, axis_wo_b, *default):
    """
    :param int axis_wo_b:
    """
    if default and axis_wo_b not in self:
      default, = default
      return default
    res = self[axis_wo_b]
    del self[axis_wo_b]
    return res

  def clear(self):
    """
    Remove all.
    """
    raise Exception("%s: cannot clear size_placeholder" % self.data)

  def keys(self):
    """
    :rtype: list[int]
    """
    return [i for i in range(self.data.ndim) if i in self]

  def values(self):
    """
    :rtype: list[tf.Tensor]
    """
    return [self[i] for i in self.keys()]

  def items(self):
    """
    :rtype: list[(int,tf.Tensor)]
    """
    return [(i, self[i]) for i in self.keys()]

  def copy(self):
    """
    :return: a copy-like object
    :rtype: dict[int,tf.Tensor]
    """
    return self.as_dict()

  def as_dict(self):
    """
    :rtype: dict[int,tf.Tensor]
    """
    return dict(self.items())


def _batch_dim_axis_from_dim_tags_tuple(dim_tags):
  """
  :param tuple[Dim] dim_tags:
  :return: batch_dim_axis. int or None if not existing
  :rtype: int|None
  """
  for axis, dim_tag in enumerate(dim_tags):
    if dim_tag.is_batch_dim():
      return axis
  return None


def _batch_shape_from_shape(shape, batch_dim_axis):
  """
  :param tuple[int|None]|list[int|None] shape: without batch-dim
  :param int|None batch_dim_axis:
  :return: shape with batch dim if existing
  :rtype: tuple[int|None]
  """
  shape = tuple(shape)
  if batch_dim_axis is not None:
    assert 0 <= batch_dim_axis <= len(shape)
    return shape[:batch_dim_axis] + (None,) + shape[batch_dim_axis:]
  else:
    return shape


# noinspection PyShadowingNames
def _create_size_placeholder(name, axis_wo_b, tag, batch_dim):
  """
  :param str name:
  :param int axis_wo_b:
  :param Dim tag:
  :param Dim|None batch_dim:
  """
  # Note on batch info: Usually, this is called early when no global batch info is initialized yet.
  # Then it is later initialized via ExternData.init_batch_info.
  # Some other external code (e.g. returnn-common) might have set custom batch info
  # on some Data instance, and via Data._adapt_batch_consistent_dim_tags / Dim.get_for_batch_ctx,
  # that might have been set on uninitialized dim tags as well.
  # Now when we get such dim tags here, they would have some Dim.batch info set
  # but this does not correspond to the global batch info which we will get here.
  # Only trust batch_dim here, or if that batch info is unset, then leave it uninitialized,
  # or even explicitly set it to None, see below.
  from .basic import reuse_name_scope
  with reuse_name_scope("extern_data/placeholders/%s" % name, absolute=True):
    dyn_size_ext = Data(
      "%s_dim%i_size" % (name, axis_wo_b), dtype=Data.size_dtype,
      dim_tags=[batch_dim] if batch_dim else [],
      batch=batch_dim.batch if batch_dim else None)
    dyn_size = tf_compat.v1.placeholder(
      name=dyn_size_ext.name, dtype=dyn_size_ext.dtype, shape=dyn_size_ext.batch_shape)
    dyn_size_ext.placeholder = dyn_size
    if dyn_size_ext.batch:
      tag.set_dyn_size_ext_for_batch_ctx(
        batch=dyn_size_ext.batch, ctx=dyn_size_ext.control_flow_ctx, dyn_size_ext=dyn_size_ext)
      # Do not set tag.batch. set_dyn_size_ext_for_batch_ctx should cover this.
    else:
      tag.batch = None  # reset, it is anyway invalid, see above
      tag.dyn_size_ext = dyn_size_ext
    tag.set_tag_on_size_tensor(dyn_size)


def _infer_dim_tags_tuple_from_shape(
  shape,
  batch_dim_axis, time_dim_axis, feature_dim_axis,
  sparse,
  size_placeholder,
  dim_tags,
  name,
  auto_create_placeholders
):
  """
  :param tuple[int|None]|list[int|None] shape: this is without batch-dim-axis
  :param int|None batch_dim_axis:
  :param int|None time_dim_axis:
  :param int|None|NotSpecified feature_dim_axis:
  :param bool sparse:
  :param dict[int,tf.Tensor]|None size_placeholder: key is axis without batch-dim
  :param dict[int,Dim]|None dim_tags: some existing explicitly specified dim tags. key is axis with batch-dim
  :param bool auto_create_placeholders:
  :param str name:
  :return: dim tags tuple
  :rtype: tuple[Dim]
  """
  assert isinstance(shape, (tuple, list))
  shape = tuple(shape)
  batch_shape = _batch_shape_from_shape(shape, batch_dim_axis=batch_dim_axis)
  if feature_dim_axis is NotSpecified:
    feature_dim_axis = _default_feature_dim_axis(
      batch_dim_axis=batch_dim_axis, time_dim_axis=time_dim_axis, batch_shape=batch_shape, sparse=sparse)
  elif feature_dim_axis is not None:
    if feature_dim_axis < 0:
      feature_dim_axis += len(batch_shape)
    assert 0 <= feature_dim_axis < len(batch_shape)
  dim_tags = dim_tags.copy() if dim_tags else {}
  if batch_dim_axis is not None and batch_dim_axis not in dim_tags:
    dim_tags[batch_dim_axis] = Dim(kind=Dim.Types.Batch, description="batch:%s" % name)
  # noinspection PyShadowingNames
  batch_dim = dim_tags[batch_dim_axis] if batch_dim_axis is not None else None
  # Note: Consistent to Data.get_dim_tag,
  # prefer interpretation as spatial axis if there is a dynamic size or this is marked as time axis.
  if size_placeholder:
    for axis_wo_b, size in size_placeholder.items():
      axis = _get_axis_wb(axis_wo_b, batch_dim_axis=batch_dim_axis)
      if axis in dim_tags:
        continue
      tag = Dim.get_tag_from_size_tensor(size)
      if tag:
        dim_tags[axis] = tag
  # See Data.get_spatial_batch_axes
  spatial_axes = [
    axis
    for axis in range(len(batch_shape))
    if axis != batch_dim_axis
    and (axis != feature_dim_axis or
         axis == time_dim_axis or
         batch_shape[axis] is None)]
  for axis in range(len(batch_shape)):
    tag = dim_tags.get(axis)
    axis_wo_b = _get_axis_wo_b(axis, batch_dim_axis=batch_dim_axis)
    dyn_size = size_placeholder.get(axis_wo_b) if (size_placeholder and axis_wo_b is not None) else None
    dim = batch_shape[axis]
    if auto_create_placeholders and dim is None and dyn_size is None and axis != batch_dim_axis:
      if not tag:
        if axis == time_dim_axis:
          tag_name = "time"
        else:
          tag_name = "spatial%i" % axis
        tag = Dim(
          description="%s:var:extern_data:%s" % (tag_name, name),
          # Spatial dim tag, even if axis == feature_dim_axis. This is to keep the old behavior.
          # This is such that Dim.is_equal behaves as before, e.g. in Data.get_common_data.
          kind=Dim.Types.Spatial)
        dim_tags[axis] = tag
      _create_size_placeholder(name=name, axis_wo_b=axis_wo_b, tag=tag, batch_dim=batch_dim)
      dyn_size = tag.dyn_size
    if tag:
      # Just some sanity checks.
      assert isinstance(tag, Dim)
      assert tag.dimension == dim
      assert tag.is_same_size_tensor(dyn_size)
      continue
    if axis == feature_dim_axis and dyn_size is None and axis != time_dim_axis:
      tag = Dim(
        kind=Dim.Types.Feature, dimension=dim, description="feature:%s" % name,
        undefined=dim is None, auto_generated=True)
    else:
      assert axis in spatial_axes
      description = "time" if axis == time_dim_axis else "spatial%i" % spatial_axes.index(axis)
      if dyn_size is not None:
        # Note: This case is uncommon/unexpected (we should have a dim-tag on the dyn_size above), so be verbose,
        # and fix such cases if possible (i.e. for all newly created dynamic size tensors, set the dim-tag).
        description += ":var:%r" % dyn_size.name
      elif dim is None:
        description += ":var-unk"
      else:
        description += ":static%i" % dim
      description += ":%s" % name
      tag = Dim(
        kind=Dim.Types.Spatial, description=description, dimension=dim, dyn_size=dyn_size,
        undefined=dim is None and dyn_size is None, auto_generated=True)
    dim_tags[axis] = tag
  assert sorted(dim_tags.keys()) == list(range(len(batch_shape)))
  return tuple(dim_tags[axis] for axis in range(len(batch_shape)))


def _auto_create_size_placeholders_on_dim_tags(name, dim_tags):
  """
  :param str name:
  :param tuple[Dim] dim_tags:
  """
  batch_dim_axis = _batch_dim_axis_from_dim_tags_tuple(dim_tags)
  batch_dim_ = dim_tags[batch_dim_axis] if batch_dim_axis is not None else None
  if batch_dim_:
    # Do this first, in case the batch dim is used elsewhere,
    # to avoid that we use some invalid batch info.
    # noinspection PyProtectedMember
    batch_dim_._validate_in_current_graph()
  for axis, tag in enumerate(dim_tags):
    # noinspection PyProtectedMember
    tag._validate_in_current_graph()
    if tag.is_batch_dim():
      continue
    if tag.dimension is not None:
      continue
    if tag.dyn_size is not None:
      continue
    axis_wo_b = _get_axis_wo_b(axis, batch_dim_axis=batch_dim_axis)
    _create_size_placeholder(name=name, axis_wo_b=axis_wo_b, tag=tag, batch_dim=batch_dim_)


def _get_axis_wo_b(axis_wb, batch_dim_axis, batch_ndim=None):
  """
  :param int axis_wb: counted with batch-dim
  :param int|None batch_dim_axis:
  :param int|None batch_ndim: only used for axis_wb < 0. might be unknown (None)
  :return: axis counted without batch-dim
  :rtype: int|None
  """
  if axis_wb < 0:
    assert batch_ndim is not None
    assert axis_wb + batch_ndim >= 0
    axis_wb += batch_ndim
    # Do this check only in this case;
    # we call this function early in construction where batch_ndim might be invalid.
    assert 0 <= axis_wb < batch_ndim
  if batch_dim_axis is None:
    return axis_wb
  if axis_wb == batch_dim_axis:
    return None
  if axis_wb < batch_dim_axis:
    return axis_wb
  return axis_wb - 1


def _get_axis_wb(axis_wo_b, batch_dim_axis):
  """
  :param int axis_wo_b: counted without batch-dim
  :param int|None batch_dim_axis:
  :return: axis counted with batch-dim
  :rtype: int
  """
  if batch_dim_axis is None:
    return axis_wo_b
  if axis_wo_b >= batch_dim_axis:
    return axis_wo_b + 1
  return axis_wo_b


def _infer_default_shape_and_time(batch_dim_axis, time_dim_axis, feature_dim_axis, sparse, dim):
  """
  This is the logic to infer some sensible/default shape when it is not specified.
  As this is somewhat adhoc, this is not recommended to be used anymore.

  :param int|None batch_dim_axis:
  :param int|None time_dim_axis:
  :param int|None|NotSpecified feature_dim_axis:
  :param bool sparse:
  :param int|None dim:
  :return: shape (without batch dim), time_dim_axis
  :rtype: (tuple[int|None],int|None)
  """
  if time_dim_axis is not None:
    assert time_dim_axis != batch_dim_axis
    shape = (None,) * (_get_axis_wo_b(time_dim_axis, batch_dim_axis=batch_dim_axis) + 1)
  else:  # no time-dim-axis
    shape = ()
  if not sparse and feature_dim_axis is not None:
    assert dim is not NotSpecified, "no shape specified, not sparse, feature_dim_axis existing -> need dim"
    if feature_dim_axis is NotSpecified or feature_dim_axis == -1:
      shape = shape + (dim,)
    else:
      assert 0 <= feature_dim_axis != batch_dim_axis
      feature_dim_axis_wo_batch = _get_axis_wo_b(feature_dim_axis, batch_dim_axis=batch_dim_axis)
      if feature_dim_axis_wo_batch < len(shape):
        shape = shape[:-feature_dim_axis_wo_batch] + (dim,) + shape[feature_dim_axis_wo_batch + 1:]
      else:
        shape = shape + (None,) * (feature_dim_axis_wo_batch - len(shape)) + (dim,)
        assert len(shape) == feature_dim_axis_wo_batch + 1
  return shape, time_dim_axis


def _default_time_dim_axis(batch_dim_axis, shape):
  """
  :param int|None batch_dim_axis:
  :param tuple[int|None]|list[int|None] shape: without batch-dim
  :return: time dim axis, counted with batch-dim
  :rtype: int|None
  """
  if batch_dim_axis is None:
    time_dim_axis = None
  else:
    # Do not select the batch dim axis, or any axis with None dim.
    # Note that we currently allow to select the same as the feature dim axis,
    # in case the feature dim is None.
    taken_axes = {batch_dim_axis}
    batch_shape = _batch_shape_from_shape(shape, batch_dim_axis=batch_dim_axis)
    for axis, _dim in enumerate(batch_shape):
      if _dim is not None:
        taken_axes.add(axis)
    available_axes = [i for i in range(len(batch_shape)) if i not in taken_axes]
    if available_axes:
      time_dim_axis = available_axes[0]
    else:
      time_dim_axis = None
  return time_dim_axis


def _default_time_dim_axis_no_shape(batch_dim_axis, feature_dim_axis):
  """
  :param int|None batch_dim_axis:
  :param int|None|NotSpecified feature_dim_axis:
  :return: time dim axis, counted with batch-dim
  :rtype: int|None
  """
  if batch_dim_axis is None:
    time_dim_axis = None
  else:
    # By default if not specified, we have a time dim.
    taken_axes = {batch_dim_axis}
    if isinstance(feature_dim_axis, int):
      taken_axes.add(feature_dim_axis)
    time_dim_axis = [i for i in range(max(taken_axes) + 2) if i not in taken_axes][0]
  return time_dim_axis


def _default_time_dim_axis_dim_tags(dim_tags):
  """
  :param list[Dim]|tuple[Dim] dim_tags:
  :return: time dim axis, counted with batch-dim
  :rtype: int|None
  """
  # Consistent to _default_time_dim_axis.
  # Any spatial dynamic axis.
  # Or otherwise any dynamic axis (including maybe feature).
  # Not using any static axes.
  dim_tags_dyn_spatial = [i for i, tag in enumerate(dim_tags) if tag.is_spatial_dim() and tag.dimension is None]
  if dim_tags_dyn_spatial:
    return dim_tags_dyn_spatial[0]
  dim_tags_dyn = [i for i, tag in enumerate(dim_tags) if not tag.is_batch_dim() and tag.dimension is None]
  if dim_tags_dyn:
    return dim_tags_dyn[0]
  return None


def _default_feature_dim_axis(batch_dim_axis, time_dim_axis, batch_shape, sparse):
  """
  :param int|None batch_dim_axis:
  :param int|None time_dim_axis:
  :param tuple[int|None] batch_shape:
  :param bool sparse:
  :return: feature dim axis, counted with batch-dim
  :rtype: int|None
  """
  if sparse:
    return None
  batch_ndim = len(batch_shape)
  ndim = batch_ndim if batch_dim_axis is None else (batch_ndim - 1)
  if ndim == 0:
    return None
  axes = [i for i in range(batch_ndim) if i not in [batch_dim_axis, time_dim_axis]]
  if not axes:
    return None
  static_axes = [i for i in axes if batch_shape[i] is not None]
  # Prefer last static, if available.
  if static_axes:
    return static_axes[-1]
  return axes[-1]


def _get_merged_dim_kind(dim_tags):
  """
  :param list[Dim]|tuple[Dim] dim_tags:
  :return: dim kind
  :rtype: Entity
  """
  if any(tag.is_batch_dim() for tag in dim_tags):
    return Dim.Types.Batch
  elif any(tag.is_feature_dim() for tag in dim_tags):
    return Dim.Types.Feature
  else:
    return Dim.Types.Spatial


class ControlFlowContext:
  """
  This is a simple wrapper around the TF ControlFlowContext, i.e. tf.while_loop or tf.cond.

  We have this wrapper to refer to a context which might not exist yet (e.g. at template construction time).
  Also, we might want to store additional information, such the spatial dim tag of the loop.
  """

  class Types:
    """
    Possible types of context.
    """
    Loop = "loop"
    Cond = "cond"

  def __init__(self, kind, outer_ctx=None):
    """
    :param str kind: from ControlFlowContext.Types
    :param ControlFlowContext outer_ctx:
    """
    self.kind = kind
    self._outer_ctx = outer_ctx
    from tensorflow.python.ops.control_flow_ops import ControlFlowContext as TFControlFlowCtx
    self._tf_control_flow_ctx = None  # type: typing.Optional[TFControlFlowCtx]
    self._loop_spatial_dim = None  # type: typing.Optional[Dim]

  def __repr__(self):
    return "ControlFlowContext{%s}" % self.repr_inner()

  def repr_inner(self):
    """
    :rtype: str
    """
    return "/".join(ctx._repr_single() for ctx in self._abs_ctx_stack())

  def _repr_single(self):
    """
    :rtype: str
    """
    s = self.kind
    if self.is_loop() and self.loop_spatial_dim:
      s += "(%s)" % self.loop_spatial_dim.short_repr()
    return s

  def _abs_ctx_stack(self):
    """
    :rtype: list[ControlFlowContext]
    :return: chain of ctx, last is self
    """
    chain = []
    ctx = self
    while ctx:
      chain.append(ctx)
      ctx = ctx.outer_ctx
    chain.reverse()
    return chain

  @classmethod
  def abs_ctx_stack(cls, ctx):
    """
    :param ControlFlowContext|None ctx:
    :rtype: list[ControlFlowContext]
    """
    if ctx:
      return ctx._abs_ctx_stack()
    return []

  @classmethod
  def abs_ctx_stack_with_root(cls, ctx):
    """
    :param ControlFlowContext|None ctx:
    :rtype: list[ControlFlowContext|None]
    :return: chain of ctx, last is self, first is None
    """
    ls = [None]  # type: typing.List[typing.Optional[ControlFlowContext]]
    if ctx:
      ls += ctx._abs_ctx_stack()
    return ls

  @classmethod
  def is_parent_or_same(cls, parent, child):
    """
    :param ControlFlowContext|None parent:
    :param ControlFlowContext|None child:
    :rtype: bool
    """
    if parent == child:
      return True
    if not parent:
      return True  # parent is root
    if not child:
      return False  # child is root but parent is not
    while child:
      if child == parent:
        return True
      child = child.outer_ctx
    return False

  @classmethod
  def collect_parent_dims(cls, ctx):
    """
    :param ControlFlowContext|None ctx:
    :rtype: list[Dim]
    """
    dims = []
    for ctx_ in ControlFlowContext.abs_ctx_stack(ctx):
      if ctx_.is_loop() and ctx_.loop_spatial_dim:
        dims.append(ctx_.loop_spatial_dim)
    return dims

  def is_loop(self):
    """
    :rtype: bool
    """
    return self.kind == self.Types.Loop

  def is_cond(self):
    """
    :rtype: bool
    """
    return self.kind == self.Types.Cond

  @property
  def outer_ctx(self):
    """
    :rtype: ControlFlowContext|None
    """
    return self._outer_ctx

  @property
  def tf_control_flow_ctx(self):
    """
    :rtype: tensorflow.python.ops.control_flow_ops.ControlFlowContext|None
    """
    return self._tf_control_flow_ctx

  @tf_control_flow_ctx.setter
  def tf_control_flow_ctx(self, ctx):
    """
    :param tensorflow.python.ops.control_flow_ops.ControlFlowContext ctx:
    """
    if self.is_loop():
      assert ctx.IsWhileContext()
    if self.is_cond():
      assert ctx.IsCondContext()
    self._tf_control_flow_ctx = ctx

  @property
  def loop_spatial_dim(self):
    """
    :rtype: Dim|None
    """
    assert self.is_loop()
    return self._loop_spatial_dim

  @loop_spatial_dim.setter
  def loop_spatial_dim(self, dim):
    """
    :param Dim dim:
    """
    assert self.is_loop()
    self._loop_spatial_dim = dim


def _dim_cmp_value(obj):
  """
  :param Dim|_MarkedDim obj:
  :return: anything which can be compared
  """
  # Make Dim and _MarkedDim comparable to each other.
  # Note that the order is really arbitrary and does not matter.
  if isinstance(obj, Dim):
    # noinspection PyProtectedMember
    return "", obj.get_same_base()._creation_idx
  if isinstance(obj, _MarkedDim):
    return obj.__class__.__name__, obj.tag
  return obj
