"""
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 returnn.util.basic as util
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.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 dyn_size_ext:
            dyn_size_ext = dyn_size_ext.copy()
        if not batch and dyn_size_ext:
            batch = dyn_size_ext.batch
            if not control_flow_ctx:
                control_flow_ctx = dyn_size_ext.control_flow_ctx
        if not batch and derived_from_tag:
            batch = derived_from_tag.batch
            if not control_flow_ctx:
                control_flow_ctx = derived_from_tag.control_flow_ctx
        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
        if derived_from_tag:
            auto_generated = derived_from_tag.auto_generated
        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
        if self.derived_from_op and self.is_dynamic():
            self.complete_dyn_size()

    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_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
        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 not self.is_dynamic():
            # 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()
        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():
                assert tag.batch == batch  # some code updated batch directly (incorrectly) and could trigger this
                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.
        if ctx:
            # Check if the ctx is really relevant, when this is derived from other tags.
            derived_bases = same_base.get_derived_bases_set()
            derived_bases.remove(same_base)
            if derived_bases:
                derived_ctxs = set()
                for d in derived_bases:
                    with util.guard_infinite_recursion(Dim.get_for_batch_ctx, d):
                        d = d.get_for_batch_ctx(batch=batch, ctx=ctx)
                    if d.control_flow_ctx:
                        derived_ctxs.add(d.control_flow_ctx)
                if not derived_ctxs:
                    ctx = None
                elif len(derived_ctxs) == 1:
                    ctx = derived_ctxs.pop()
                else:
                    raise NotImplementedError("not yet implemented: multiple derived ctxs: %r" % (derived_ctxs,))
        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 and not same_base.derived_from_op:
            return None
        if not dyn_size_ext and same_base._same_for_batch_ctx:
            # There are earlier entries in _same_for_batch_ctx
            # -- maybe we can infer dyn_size_ext, even with different batch.
            for (batch_, ctx_), other in same_base._same_for_batch_ctx.items():
                if ctx_ == ctx and other.dyn_size_ext:
                    dyn_size_ext = other.dyn_size_ext.copy_template()
                    dyn_size_ext.beam = batch.beam
                    dyn_size_ext.batch = batch
                    break
        ctx = dyn_size_ext.control_flow_ctx if dyn_size_ext else ctx
        dim_tag = None
        for candidate in [self, same_base]:
            if (
                (candidate.batch == batch or (not candidate.batch and batch.is_global_batch()))
                and not candidate.control_flow_ctx
                and not ctx
            ):
                # The same_base instance is either undefined (no batch, no ctx)
                # or it is defined for the same batch and ctx.
                # In any case, reuse it then.
                candidate.batch = batch
                if dyn_size_ext:
                    if candidate.dyn_size_ext:
                        candidate.dyn_size_ext.batch = batch
                        assert candidate.dyn_size_ext.dim_tags == dyn_size_ext.dim_tags
                    else:
                        candidate.dyn_size_ext = dyn_size_ext
                    assert not candidate.dyn_size_ext.control_flow_ctx
                elif candidate.dyn_size_ext:
                    candidate.dyn_size_ext.batch = batch
                else:
                    candidate.complete_dyn_size(template_only=True)
                dim_tag = candidate
                break
        if not dim_tag:
            dim_tag = Dim(
                kind=self.kind,
                description=self.description,
                dimension=self.dimension,
                auto_generated=self.auto_generated,
                batch=batch,
                control_flow_ctx=ctx,
                dyn_size_ext=dyn_size_ext,
            )
            dim_tag.same_as = same_base
        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[(batch, ctx)] = dim_tag
        dim_tag.complete_dyn_size(template_only=True)
        return dim_tag

    def reset_batch_ctx(self):
        """
        For the self instance, reset batch and context.
        """
        self._same_for_batch_ctx.pop((self.batch, self.control_flow_ctx), None)
        self.batch = None
        self.control_flow_ctx = None
        if self.dyn_size_ext and self.dyn_size_ext.batch:
            self.dyn_size_ext = self.dyn_size_ext.copy_template()
            self.dyn_size_ext.batch = None

    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)
        assert dyn_size_ext.batch == batch and dyn_size_ext.control_flow_ctx == ctx
        if same.dyn_size_ext:
            assert same.dyn_size_ext.dim_tags == dyn_size_ext.dim_tags
            if dyn_size_ext.placeholder is not None:
                same.dyn_size_ext.placeholder = dyn_size_ext.placeholder
        else:
            same.dyn_size_ext = dyn_size_ext
        self._maybe_update()

    def get_dyn_size_ext_for_batch_ctx(self, batch, ctx, template_only=False):
        """
        :param BatchInfo|None batch:
        :param ControlFlowContext|None ctx:
        :param bool template_only:
        :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
        same.complete_dyn_size(template_only=template_only)
        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.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:
        """
        if self.dyn_size_ext and self.dyn_size_ext.placeholder is dyn_size:  # fast path check
            return
        assert self.can_be_used_as_dim()
        assert isinstance(dyn_size, tf.Tensor) and dyn_size.shape.ndims == 1
        other = Dim.get_tag_from_size_tensor(dyn_size)
        if other:
            self.declare_same_as(other)
            if self.batch:
                assert self.batch == other.batch and self.control_flow_ctx == other.control_flow_ctx
            else:
                self.batch = other.batch
                self.control_flow_ctx = other.control_flow_ctx
            self.dyn_size_ext = other.dyn_size_ext
            return
        self._init_default_dyn_size_ext(dyn_size)
        self.set_tag_on_size_tensor(dyn_size)

    def _init_default_dyn_size_ext(self, dyn_size):
        """
        :param tf.Tensor dyn_size:
        """
        if self.dyn_size_ext:
            if self.dyn_size_ext.placeholder is not None:
                # Do not allow resetting it to sth different.
                assert self.dyn_size_ext.placeholder is dyn_size
        else:
            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,
                shape=(),
                batch_dim_axis=0,
                batch=self.batch,
                beam=beam,
                control_flow_ctx=self.control_flow_ctx,
            )
        self.dyn_size_ext.placeholder = 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 not self.is_dynamic() and 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 None and not self.is_batch_dim()

    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
        tag = Dim.get_tag_from_size_tensor(x)
        if tag and tag == self:
            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. ",
                            "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"
            ), "%s: dyn size %s has unexpected batch %s, expected %s" % (
                self,
                x,
                batch,
                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 not self.is_batch_dim() and self.is_dynamic():
            if same_as_before:
                assert self.dyn_size_ext and self.dyn_size_ext.placeholder is not None
                # Do not overwrite it.
            else:
                self._init_default_dyn_size_ext(x)
        if getattr(x, "_is_size_of_dim_tag", None) is None:
            setattr(x, "_is_size_of_dim_tag", self)
        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, template_only=False):
        """
        In case we can calculate the dyn size, do that now.

        :param bool template_only:
        """
        if not self.is_dynamic():
            return
        self._validate_in_current_graph()
        if self.dyn_size_ext and (self.dyn_size_ext.placeholder is not None or template_only):
            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 template_only:
                    return None
                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 not x.is_dynamic():  # static
                    assert 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=None if template_only else 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(template_only=template_only)
                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
        """
        from returnn.util import BehaviorVersion

        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 BehaviorVersion.get() >= 16:
            # Either self or other is some dim tag explicitly created by the user,
            # and they are not the same, so we never treat them as equal.
            if not self.auto_generated or not other.auto_generated:
                if broadcast_matches and (
                    (self.dimension == 1 and self.auto_generated) or (other.dimension == 1 and other.auto_generated)
                ):
                    pass  # exception, allow broadcast logic
                else:
                    return False
        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(())
        with util.guard_infinite_recursion(Dim.__hash__, self):
            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_set(self):
        """
        :rtype: set[Dim]
        """
        res = set()
        queue = [self]
        visited = {}  # type: typing.Dict[int,Dim]  # by id
        while queue:
            base = queue.pop(-1)
            if base.same_as:
                base = base.same_as
            if id(base) in visited:
                continue
            visited[id(base)] = base
            res.add(base)
            if base.derived_from_op:
                queue.extend(base.derived_from_op.inputs)
            elif base.derived_from_tag:
                queue.append(base.derived_from_tag)
        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
        # Note: Check `is`, not `==`. `==` can be true even though same_as are not the same instance,
        # e.g. via auto_generated.
        if self is other:
            return
        self._maybe_update()
        self._validate_in_current_graph()
        other._validate_in_current_graph()
        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.
            with util.guard_infinite_recursion(Dim.declare_same_as, other, self):
                return other.declare_same_as(self)
        if self.batch:
            # If self is defined (self.is_dim_known), be fair to other, and adapt it to the right batch,
            # such that other.is_dim_known is correct, by potentially completing it.
            other_ = other.get_for_batch_ctx(self.batch, ctx=self.control_flow_ctx)
        else:
            other_ = other
        if (
            (self.is_dim_known() and not other_.is_dim_known())
            or
            # Like is_dim_known but for static dims, we might know both,
            # but the derived_from_op still would provide more information.
            (
                self_same_as.derived_from_op
                and not other_same_base.derived_from_op
                and other not in self.get_derived_bases_set()
            )
            or (not self.undefined and other_.undefined)
        ):
            if self_same_as._creation_idx > other_same_base._creation_idx:
                # We want to keep other instead.
                # https://github.com/rwth-i6/returnn_common/issues/200
                self_same_as.description = other_same_base.description
                # Set new description, because in case of auto_generated,
                # they would evaluate to be the same dim otherwise.
                other_same_base.description = other_same_base.description + "{to-be-replaced}"
            with util.guard_infinite_recursion(Dim.declare_same_as, other, self):
                return other.declare_same_as(self)
        other_derived_bases = other.get_derived_bases_set()
        self_derived_bases = self.get_derived_bases_set()
        if other_derived_bases != self_derived_bases and self_derived_bases.issubset(other_derived_bases):
            # Avoid cycles on derived_from_tag. https://github.com/rwth-i6/returnn/issues/1054
            with util.guard_infinite_recursion(Dim.declare_same_as, other, self):
                return other.declare_same_as(self)
        self.derived_from_op = None
        self.derived_from_tag = None
        if self_same_as is not self:
            assert not self_same_as.same_as
            if self_same_as is other_same_base:
                return
            with util.guard_infinite_recursion(Dim.declare_same_as, self_same_as, other_same_base):
                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, template_only=True
                )
        else:
            for dim_ in self._same_for_batch_ctx.values():
                dim_.derived_from_op = None
                dim_.derived_from_tag = None
        # Now merge existing variants. But only if not derived via op, because in that case, we can (and should!)
        # automatically infer it. Note that we only got here when the other is not the same dim, so it means that
        # the other is really different, the sizes are potentially different, but we want to overtake the other.
        if other_same_base.derived_from_op:
            # Cleanup everything, esp potential already computed sizes, as these might be invalid.
            for dim_ in [self_same_as, self] + list(self._same_for_batch_ctx.values()):
                if dim_.dyn_size_ext:
                    dim_.dyn_size_ext.placeholder = None
        other_same_base._merge_same_for_batch_ctx_dict(self)
        other._maybe_update()
        self.same_as = other_same_base
        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, template_only=True
                )
                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, template_only=True
                )
                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, template_only=True
            )
            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, template_only=True
            )
            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
        self.auto_generated = self_same_as.auto_generated = other_same_base.auto_generated
        # Take over derived_from_op. However, only if this would not introduce cycles!
        if not self_derived_bases.issuperset(other_derived_bases):
            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
        if self.batch:
            self_ = self.get_for_batch_ctx(batch=self.batch, ctx=self.control_flow_ctx)
            if self_ is not self:
                self.control_flow_ctx = self_.control_flow_ctx  # might be different
                self.dyn_size_ext = self_.dyn_size_ext  # might be unset

    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

    def derive_from(self, base, set_derived_from_flag=True):
        """
        :param Dim base:
        :param bool set_derived_from_flag:
        """
        self_base = self.get_same_base()
        if set_derived_from_flag:
            if self_base.derived_from_tag:
                assert self_base.derived_from_tag == base
            else:
                self_base.derived_from_tag = base
        if not self.batch and base.batch:
            self.batch = base.batch
            self.control_flow_ctx = base.control_flow_ctx
            key = base.batch, base.control_flow_ctx
            assert key not in self_base._same_for_batch_ctx
            self_base._same_for_batch_ctx[key] = self
        if self.is_dynamic() or not self.is_dim_known():
            if not self.dyn_size_ext:
                if base.dyn_size_ext:
                    if base.batch and base.batch == self.batch and base.control_flow_ctx == self.control_flow_ctx:
                        self.dyn_size_ext = base.dyn_size_ext.copy_template(name="%s:size" % self_base.description)
                elif base.is_batch_dim():
                    self.dyn_size_ext = Data(
                        "%s:batch" % self_base.description, shape=(), dtype="int32", batch_dim_axis=None
                    )

    @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_set():
                        replace_existing = True
                    if replace_existing:  # Replace the existing by the new tag.
                        tags[tags.index(existing_tag)] = tag
                        for _, dims_ in data_axes_dict.items():
                            dims_[:] = [tag if d == existing_tag else d for d in dims_]
                        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)),
                derived_from_tag=self.representative_tag(),
            )

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

    @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_.is_dynamic():
                        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)
            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,
                extern_data=auto_create_placeholders,
                dim_tags=dim_tags,
                sparse=sparse,
                batch=self.batch,
            )
            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)
                self._dim_tags = self._dim_tags[:_axis] + (_dim_tag,) + self._dim_tags[_axis + 1 :]
        if auto_create_placeholders:
            # Do that after same_dim_tags_as handling.
            _auto_create_size_placeholders_on_dim_tags(name=name, dim_tags=self._dim_tags)
        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
        if assume_complete:
            for tag in self.dim_tags:
                if tag.is_batch_dim():
                    continue
                if tag.is_dynamic():
                    assert tag.dyn_size_ext, "%s sanity_check: dynamic dim %s undefined" % (self, tag)
                    if not ignore_placeholder:
                        if tag.dyn_size_ext.placeholder is None:
                            tag.complete_dyn_size()
                        if self.placeholder is not None:
                            assert (
                                tag.dyn_size_ext.placeholder is not None
                            ), "%s sanity_check: dynamic dim %s value unknown" % (self, tag)
                assert tag.is_dim_known()

    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, allow_missing_implicit_dims=False):
        """
        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`.
        :param bool allow_missing_implicit_dims:
        """
        actual_dims_str = "{%s}" % ", ".join(
            [str(d) for d in list(self.dim_tags) + sorted(self.dim_tags_set_implicit_only_wrapped)]
        )  # noqa
        expected_dims_str = "{%s}" % ", ".join([str(d) for d in sorted(out_shape)])
        self_dim_tags = self.dim_tags_set_implicit
        self_dim_tags_implicit_only = self.dim_tags_set_implicit_only_wrapped
        if not out_shape:  # allow also empty list or empty tuple
            if self_dim_tags:
                raise VerifyOutShapeException(
                    "%s verify_out_shape:\n" % self
                    + "Actual dims: %s\nExpected empty out_shape: %s" % (actual_dims_str, expected_dims_str)
                )
            return
        if not isinstance(out_shape, set):
            # out_shape is not empty (tested above), so must be a 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:\n"
                            "Actual dims: %s\n"
                            "Expected out_shape: %s\n"
                            "%s is not an implicit dim in self"
                        )
                        % (self, actual_dims_str, expected_dims_str, 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 does not match:\n" % self
                        + "Actual dims: %s\nExpected out_shape: %s\n" % (actual_dims_str, expected_dims_str)
                        + "Dim %s multiple times in out_shape" % dim
                    )
                raise VerifyOutShapeException(
                    "%s verify_out_shape:\n" % self
                    + "Actual dims: %s\nExpected out_shape: %s\n" % (actual_dims_str, expected_dims_str)
                    + "Dim %s not in self" % dim
                )
            remaining.discard(dim_tag)
        if remaining:
            if allow_missing_implicit_dims and remaining.issubset(self.dim_tags_set_implicit_only):
                pass  # ok
            else:
                raise VerifyOutShapeException(
                    "%s verify_out_shape missing dims:\n" % self
                    + "Actual dims: %s\nExpected out_shape: %s\n" % (actual_dims_str, expected_dims_str)
                    + "Missing dims: %s" % ", ".join(map(str, sorted(remaining)))
                )

    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
        if self._batch == batch:  # fast path
            return
        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]]
            elif axes.startswith("tag:"):  # any tag
                return [self.get_axis_by_tag_name(axes[len("tag:") :])]
            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.is_dynamic()
        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]]
                shape[self.batch_dim_axis] = self.get_batch_dim()
                shape[axis] = tag.get_dim_value()
                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_dyn_size_tags(self):
        """
        :return: all dim tags with dynamic size
        :rtype: list[Dim]
        """
        return [dim_tag for dim_tag in self._dim_tags if dim_tag.is_dynamic()]

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

    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, tag %s" % (
                other_axes,
                other,
                self,
                other_axis,
                other.dim_tags[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.reset_batch_ctx()  # 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, batch, size_placeholder, dim_tags, name, extern_data
):
    """
    :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 BatchInfo|None batch:
    :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 extern_data:
    :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)
    # 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 extern_data 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,
                    batch=batch,
                    auto_generated=True,
                )
                tag.dyn_size_ext = Data(
                    "%s_dim%i_size" % (name, axis_wo_b), dtype=Data.size_dtype, shape=(), batch=batch
                )
                dim_tags[axis] = tag
            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,
                batch=batch if dim is None else None,
                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,
                batch=batch if dim is None else None,
                undefined=dim is None and dyn_size is None,
                auto_generated=True,
            )
        if dim is None and tag.dyn_size_ext is None:
            tag.dyn_size_ext = Data("%s_dim%i_size" % (name, axis_wo_b), dtype=Data.size_dtype, shape=(), batch=batch)
            if dyn_size is not None:
                tag.dyn_size_ext.placeholder = dyn_size
        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 not tag.is_dynamic():
            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:
            try:
                with util.guard_infinite_recursion(ControlFlowContext._repr_single, self):
                    s += "(%s)" % self.loop_spatial_dim.short_repr()
            except util.InfiniteRecursionDetected as exc:
                # repr should always return and not throw errors
                s += "(%s for loop_spatial_dim?)" % exc
        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
