"""
Backwards-compatible functions and attribs for the old ``Dim`` class,
or just rarely used attribs, such that we can save memory for the common case.
"""

from __future__ import annotations
from typing import TYPE_CHECKING, Optional, Union, Tuple, Sequence, Dict, List

from returnn.util.basic import Entity
from returnn.util import basic as util

if TYPE_CHECKING:
    # Those are only used for TensorFlow, or they are deprecated.
    from returnn.tf.util.data import BatchInfo, ControlFlowContext

    # just for type hints, otherwise use _d.Dim
    from .dim import Dim

from . import dim as _d
from . import tensor as _t
from . import marked_dim as _m


class DimTypes:
    """
    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)


class _DimExtra:
    def __init__(
        self,
        *,
        dim: Dim,
        kind=DimTypes.Unspecified,
        vocab=None,
        undefined=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: Optional[_t.Tensor] = None,
        src_axis: Optional[int] = None,
    ):
        """
        :param dim:
        :param Entity|None kind:
        :param returnn.datasets.util.vocabulary.Vocabulary|None vocab:
        :param bool undefined: When this is specified as `None` by the user via `shape`.
        :param bool special: this can not be a dim tag of :class:`Tensor`.
            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 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 src_data:
        :param src_axis:
        """
        self.dim = dim
        assert kind is None or (isinstance(kind, Entity) and kind in DimTypes.Types)
        self.kind = kind
        self.vocab = vocab
        self.same_as = None  # type: Optional[_d.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, _t.Tensor) and isinstance(src_axis, int)
        if not batch and dim.dyn_size_ext:
            batch = dim.dyn_size_ext.batch
            if not control_flow_ctx:
                control_flow_ctx = dim.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
        self.dyn_size_same = set()  # set of TensorRef
        self.undefined = undefined
        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: Dict[Tuple[BatchInfo,Optional[ControlFlowContext]],_d.Dim]

    def __getstate__(self):
        d = vars(self).copy()
        d["batch"] = None
        d["same_for_batch_ctx"] = {}
        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 DimTypes.Types}[self.kind]


class _DimMixin:
    name: Optional[str]
    capacity: Optional[int]
    size: Optional[int]
    dyn_size_ext: Optional[_t.Tensor]
    _extra: Optional[_DimExtra]

    def _handle_extra_kwargs(self: Dim, *, dyn_size: Optional[_t.RawTensorType] = None, **kwargs):
        if kwargs:
            self._extra = _DimExtra(dim=self, **kwargs)
        if dyn_size is not None:
            self.dyn_size = dyn_size
        if self.derived_from_op and self.is_dynamic():
            self.complete_dyn_size()

    @property
    def description(self) -> Optional[str]:
        """
        :return: description, alias for name
        """
        return self.name

    @property
    def dimension(self) -> Optional[int]:
        """
        :return: alias for static size, or None if dynamic
            In __init__, it is more flexible, but we require this API for the attrib (property)
            for compatibility to old code.
        """
        return self.size

    @property
    def kind(self) -> Optional[Entity]:
        """
        :return: one in DimTypes (deprecated)
        """
        if not self._extra:
            return None
        return self._extra.kind

    @property
    def match_priority(self) -> int:
        """
        :return: match priority
        """
        if not self._extra:
            return 0
        return self._extra.match_priority

    @property
    def batch(self) -> Optional[BatchInfo]:
        """
        :return: batch info (deprecated)
        """
        if not self._extra:
            if self.dyn_size_ext:
                return self.dyn_size_ext.batch
            return None
        return self._extra.batch

    @batch.setter
    def batch(self: Dim, value: Optional[BatchInfo]):
        if self.batch is value:
            return
        self._make_extra().batch = value

    @property
    def control_flow_ctx(self) -> Optional[ControlFlowContext]:
        """
        :return: control flow context (deprecated)
        """
        if not self._extra:
            if self.dyn_size_ext:
                return self.dyn_size_ext.control_flow_ctx
            return None
        return self._extra.control_flow_ctx

    @control_flow_ctx.setter
    def control_flow_ctx(self: Dim, value: Optional[ControlFlowContext]):
        if self.control_flow_ctx is value:
            return
        self._make_extra().control_flow_ctx = value

    @property
    def auto_generated(self) -> bool:
        """
        :return: see _DimExtra
        """
        if not self._extra:
            return False
        return self._extra.auto_generated

    @property
    def same_as(self) -> Optional[Dim]:
        """
        :return: same as other dim
        """
        if not self._extra:
            return None
        return self._extra.same_as

    @same_as.setter
    def same_as(self: Dim, value: Optional[_d.Dim]):
        if self.same_as is value:
            return
        self._make_extra().same_as = value

    @property
    def special(self) -> bool:
        """
        :return: see _DimExtra
        """
        if not self._extra:
            return False
        return self._extra.special

    @property
    def derived_from_op(self) -> Optional[Op]:
        """
        :return: op
        """
        if not self._extra:
            return None
        return self._extra.derived_from_op

    @property
    def derived_from_tag(self) -> Optional[Dim]:
        """
        :return: dim
        """
        if not self._extra:
            return None
        return self._extra.derived_from_tag

    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 += f"({self.dimension})"
        else:
            if self.dyn_size_ext:
                desc += "[%s]" % ",".join(self.dyn_size_ext.get_batch_axes_short_description(special_axes=False))
            else:
                desc += "[?]"
            if self.control_flow_ctx:
                desc += "{ctx=%s}" % self.control_flow_ctx.repr_inner()
        return desc

    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 __reduce_ex__(self: _d.Dim, protocol):
        if self == _d.batch_dim:
            return "batch_dim"
        if self == _d.single_step_dim:
            return "single_step_dim"
        return super().__reduce_ex__(protocol)

    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 = _d.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._extra.src_data if self._extra else None,
            src_axis=self._extra.src_axis if self._extra else None,
        )
        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
        from returnn.tf.util.data import ControlFlowContext

        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
        """
        if (
            self.dyn_size_ext and not self.dyn_size_ext.is_valid_in_current_graph()
        ):  # maybe from an earlier run which reuses the dim tag
            # Reset and cleanup.
            self.dyn_size_ext = self.dyn_size_ext.copy_template()
            self.dyn_size_ext.batch = None
            same_base = self.get_same_base()
            # noinspection PyProtectedMember
            if same_base._extra:
                # noinspection PyProtectedMember
                same_base._extra.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: Dim):
        if self.is_batch_dim():
            return
        if isinstance(self.size, int):
            return
        if not self._extra:
            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
        extra = self._get_same_base_extra()
        if not extra:
            return
        key = (self.batch, self.control_flow_ctx)
        if self.dyn_size_ext and key not in extra.same_for_batch_ctx:
            extra.same_for_batch_ctx[key] = self
        # Check if we can find more
        if key in extra.same_for_batch_ctx:
            same = extra.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: Dim, batch, ctx, allow_none=False) -> Optional[Dim]:
        """
        Warning: This is only for TensorFlow, and also we might want to remove it.
        https://github.com/rwth-i6/returnn/issues/975

        :param BatchInfo batch:
        :param ControlFlowContext|None ctx:
        :param bool allow_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 _d.Dim(kind=DimTypes.Batch, description="batch:%s" % batch.short_repr(), batch=batch, dimension=None)
        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.
        if self._extra:
            same_base = self.get_same_base()
            same_base._validate_in_current_graph()
            # noinspection PyProtectedMember
            if same_base._extra:
                from returnn.tf.util.data import ControlFlowContext

                for ctx_ in ControlFlowContext.abs_ctx_stack_with_root(ctx):
                    # noinspection PyProtectedMember
                    tag = same_base._extra.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
        else:
            same_base = self
        same_base_extra = same_base._make_extra()
        # 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(_d.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  # type: Optional[_d.Dim]
            # noinspection PyProtectedMember
            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
            elif same_base._extra:
                from returnn.tf.util.data import ControlFlowContext

                for ctx_ in ControlFlowContext.abs_ctx_stack_with_root(ctx):
                    # noinspection PyProtectedMember
                    tag = same_base._extra.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.
                        import tensorflow as tf
                        from returnn.tf.util.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_extra.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_extra.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 = _d.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 _d.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_extra.same_for_batch_ctx[(batch, ctx)] = dim_tag
        dim_tag.complete_dyn_size(template_only=True)
        return dim_tag

    def reset_batch_ctx(self: Dim):
        """
        For the self instance, reset batch and context.
        """
        if self._extra:
            self._extra.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()
        other = _d.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 = _t.Tensor(
                name=("%s:dyn_size" % self.description) if self.description else dyn_size.op.name,
                dtype=_t.Tensor.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 get_mask(self: Dim, *, dim_order: Optional[Sequence[Dim]] = None) -> _t.Tensor:
        """
        :param dim_order: if given, the dims of the mask will be in this order.
            This can be useful if the mask is broadcasted against some other tensor.
        :return: if need_masking(), the corresponding mask.
            If this is e.g. the time-dim T of shape [B], then the mask will be of shape [B,T].
            The mask could be used with :func:`masked_select` (``boolean_mask``) or ``where``.
        """
        import returnn.frontend as rf

        assert self.dyn_size_ext and self.dyn_size_ext.raw_tensor is not None
        # noinspection PyProtectedMember
        backend = self.dyn_size_ext._raw_backend

        max_idx = rf.reduce(
            self.dyn_size_ext,
            axis=self.dyn_size_ext.dims,
            mode="max",
            # Masking here is not always possible, e.g. if we have
            # tag = Dim{'self-att-keys'['time:var:extern_data:classes'[B]]}
            use_mask=False,
        )
        # We use the assumption that self.placeholder.shape[axis] == max_idx.
        # 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 = self.dyn_size_ext.copy_masked(max_idx)
        idx_range = backend.range_over_dim(self)
        seq_mask = rf.compare(idx_range, "<", size_ext, allow_broadcast_all_sources=True, dim_order=dim_order)
        return seq_mask

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

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

    def is_spatial_dim(self):
        """
        :return: whether this dim tag is of kind spatial
        :rtype: bool
        """
        return self.kind == DimTypes.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
        extra = self._get_same_base_extra()
        if extra:
            for _, other in extra.same_for_batch_ctx.items():
                if other.dyn_size_ext:
                    return True
        return False

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

    def is_static(self) -> bool:
        """
        :return: static
        """
        return not self.is_dynamic()

    def need_masking(self):
        """
        :return: whether dim is static or dynamic but with scalar dyn_size_ext
        """
        if self.is_static():
            if self.capacity is not None:
                return self.size < self.capacity
            return False
        if self.capacity is not None:
            return True
        if not self.dyn_size_ext:
            return True  # unknown
        return self.dyn_size_ext.batch_ndim > 0

    def can_be_used_as_dim(self):
        """
        :return: whether this can be used as a dim in :class:`Data`, i.e. it is not special
        :rtype: bool
        """
        return 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 self.dyn_size_ext and x is self.dyn_size_ext.placeholder:
            return True
        tag = _DimMixin.get_tag_from_size_tensor(x)
        if tag and tag == self:
            return True
        if not self._extra:
            return False
        if util.RefIdEq(x) in self._extra.dyn_size_same:
            return True
        return False

    def set_tag_on_size_tensor(self: Dim, x, batch=None, same_as_before=False) -> Dim:
        """
        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 x: raw tensor, for example tf.Tensor
        :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
        """
        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 self._extra and util.RefIdEq(x) in self._extra.dyn_size_same:
                pass  # ok, pass on
            elif same_as_before:
                self._make_extra().dyn_size_same.add(util.RefIdEq(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.
                # noinspection PyProtectedMember
                from returnn.frontend._backend import get_backend_by_raw_tensor_type

                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:",
                            get_backend_by_raw_tensor_type(type(x)).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) -> Optional[_d.Dim]:
        """
        :param tf.Tensor x: size tensor. has been set before via :func:`set_tag_on_size_tensor`
        """
        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 not op:
            return

        for x in op.inputs:
            if self.batch:
                x = x.get_for_batch_ctx(self.batch, self.control_flow_ctx)
            x.complete_dyn_size(template_only=template_only)

        backend = None
        for x in op.inputs:
            if self.batch:
                x = x.get_for_batch_ctx(self.batch, self.control_flow_ctx)
            if x.dyn_size_ext and x.dyn_size_ext.raw_tensor is not None:
                # noinspection PyProtectedMember
                backend = x.dyn_size_ext._raw_backend
                break

        size_dtype = None
        for x in op.inputs:
            if self.batch:
                x = x.get_for_batch_ctx(self.batch, self.control_flow_ctx)
            if x.dyn_size_ext:
                size_dtype = x.dyn_size_ext.dtype
                break
        if not size_dtype:
            size_dtype = _t.Tensor.size_dtype

        import numpy
        import returnn.frontend as rf
        from returnn.tensor import Tensor

        tf = tf_util = tensor_util = None
        if backend and backend.is_tensorflow:
            import tensorflow as tf

            if backend.RawTensorType == tf.Tensor:
                from returnn.tf.util import basic as tf_util
                from tensorflow.python.framework import tensor_util
            else:
                tf = None

        kind = op.kind
        if kind.endswith("_right"):
            kind = kind[: -len("_right")]  # order does not matter here
        if kind.endswith("_left"):
            kind = kind[: -len("_left")]

        def _is_negative(x__):
            if isinstance(x__, numpy.ndarray):
                return (x__ < 0).any()
            if isinstance(x__, (int, float, numpy.number)):
                return x__ < 0
            if not tf:
                return False
            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_tf(a, b):
            if template_only:
                return None
            if a is None or b is None:
                return None
            assert isinstance(a, tf.Tensor) and isinstance(b, (int, tf.Tensor))
            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)

        def _bin_op(a, b):
            if template_only or not backend:
                if isinstance(a, _t.Tensor) and isinstance(b, _t.Tensor):
                    return _t.Tensor.get_common_data([a, b], allow_broadcast_all_sources=True)
                if isinstance(a, _t.Tensor):
                    return a
                if isinstance(b, _t.Tensor):
                    return b
            if kind == "add":
                return a + b
            elif kind == "sub":
                return 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":
                if isinstance(a, Tensor):
                    return rf.ceil_divide(a, b)
                return -(-a // b)
            else:
                raise ValueError("unknown op kind %r" % op.kind)

        y_name = self.description + ":seq-length"
        y: Optional[_t.Tensor] = None  # resulting dyn size
        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:
                    if not template_only and backend and not tf:
                        y = backend.convert_to_tensor(x.dimension, dims=[], dtype=size_dtype, name=y_name)
                    else:
                        y = _t.Tensor(
                            name=y_name,
                            dim_tags=[],
                            dtype=size_dtype,
                        )
                        if not template_only and tf:
                            with tf.control_dependencies(None):  # this will reset the context
                                y.raw_tensor = tf.constant(x.dimension)
                    continue
                if tf:
                    y.placeholder = _bin_op_tf(y.placeholder, x.dimension)
                else:
                    y = _bin_op(y, 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 = _t.Tensor.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
            if tf:
                y.placeholder = _bin_op_tf(y_.placeholder, x_.placeholder)
            else:
                y = _bin_op(y_, x_)
        assert y, f"op {op}?"
        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 tf and 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 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 == DimTypes.Feature and ignore_feature_dim:
            return True
        if treat_feature_as_spatial:
            if self_kind == DimTypes.Feature:
                self_kind = DimTypes.Spatial
            if other_kind == DimTypes.Feature:
                other_kind = DimTypes.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 == DimTypes.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 == DimTypes.Feature:
            if allow_same_feature_dim:
                return True
        if self_kind == other_kind == DimTypes.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, _d.Dim):
            return False
        return self.is_equal(other)

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

    def __hash__(self):
        """
        :rtype: int
        :return: hash, matching to :func:`__eq__`
        """
        # 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.special:
            return hash(id(self))
        if self.is_batch_dim():
            return hash(())
        with util.guard_infinite_recursion(_d.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: Dim, other: Dim):
        """
        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, (_d.Dim, _m.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: _d.Dim) -> _d.Dim:
        """
        :return: same base
        """
        if not self._extra:
            return self
        base = self
        while base.same_as:
            base = base.same_as
        return base

    def get_same_derived_base(self: _d.Dim) -> _d.Dim:
        """
        :return: same base, but also consider derived_from_...
        """
        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: Dict[int,_d.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: _d.Dim) -> bool:
        """
        :return: whether the undefined flag is set, in self, bases, or any derived bases. also see :func:`is_dim_known`
        """
        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
            # noinspection PyProtectedMember
            if base._extra and base._extra.undefined:
                return True
            if base.same_as:
                base = base.same_as
                continue
            base = base.derived_from_tag
            assert base
        # noinspection PyProtectedMember
        return base._extra and base._extra.undefined

    def declare_same_as(self: _d.Dim, other: _d.Dim):
        """
        :param 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(_d.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)
        ):
            with util.guard_infinite_recursion(_d.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(_d.Dim.declare_same_as, other, self):
                return other.declare_same_as(self)
        if self._extra:
            self._extra.derived_from_op = None
            self._extra.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(_d.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:
            if self._extra:
                for dim_ in self._extra.same_for_batch_ctx.values():
                    # noinspection PyProtectedMember
                    if dim_._extra:
                        # noinspection PyProtectedMember
                        dim_._extra.derived_from_op = None
                        # noinspection PyProtectedMember
                        dim_._extra.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._extra.same_for_batch_ctx.values()) if self._extra else []):
                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._extra and self._extra.src_data:
            assert isinstance(self._extra.src_axis, int)
            # Maybe it changed in the meanwhile, so check.
            tag = self._extra.src_data.get_dim_tag(self._extra.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.capacity = self.capacity
            other.size = self.size
        elif not self.is_dim_known() and other.is_dim_known():
            self.capacity = other.capacity
            self.size = other.size
        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._make_extra()
        self_same_as._make_extra()
        # noinspection PyProtectedMember
        self._extra.auto_generated = self_same_as._extra.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:
                # noinspection PyProtectedMember
                other_same_base._make_extra().derived_from_op = self.derived_from_op
            elif other_same_base.derived_from_op and not self.derived_from_op:
                self._make_extra().derived_from_op = other_same_base.derived_from_op
        if self._extra and not other_same_base.is_dynamic():
            # Those might be set via get_batch_for_ctx for an undefined dim,
            # which now becomes static due to `other`.
            self._extra.batch = None
            self._extra.control_flow_ctx = None
            for key, dim_ in self._extra.same_for_batch_ctx.items():
                # noinspection PyProtectedMember
                dim_extra = dim_._extra
                if dim_extra:
                    dim_extra.batch = None
                    dim_extra.control_flow_ctx = None
        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: _d.Dim, other: _d.Dim):
        """
        :param other:
        """
        # noinspection PyProtectedMember
        if not self._extra and not other._extra:
            return
        self._validate_in_current_graph()
        if self._extra:
            for _, dim in list(self._extra.same_for_batch_ctx.items()):
                assert isinstance(dim, _d.Dim)
                dim._validate_in_current_graph()
        # noinspection PyProtectedMember
        if other._extra:
            # noinspection PyProtectedMember
            for key, dim in other._extra.same_for_batch_ctx.items():
                if not dim._validate_in_current_graph():
                    continue
                self_dim = self._make_extra().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._extra.same_for_batch_ctx[key] = dim
            # noinspection PyProtectedMember
            other._extra.same_for_batch_ctx.clear()  # we only want to have it once

    # noinspection PyProtectedMember
    def derive_from(self: _d.Dim, base: _d.Dim, set_derived_from_flag: bool = True):
        """
        :param base: dim
        :param set_derived_from_flag:
        """
        self_base = self.get_same_base()
        self_base_extra = self_base._make_extra()
        if set_derived_from_flag:
            if self_base_extra.derived_from_tag:
                assert self_base_extra.derived_from_tag == base
            else:
                self_base_extra.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_extra.same_for_batch_ctx
            self_base_extra.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 = _t.Tensor(
                        name="%s:batch" % self_base.description, shape=(), dtype="int32", batch_dim_axis=None
                    )

    def copy_from(self: Dim, other: Dim):
        """define"""
        self.size = other.size
        self.capacity = other.capacity
        self.dyn_size_ext = other.dyn_size_ext
        self.derive_from(other)

    @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[_t.Tensor] 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], util.DictRefKeys[_t.Tensor, list[Dim]])
        """
        tags = []
        data_axes_dict = util.DictRefKeys()  # type: util.DictRefKeys[_t.Tensor, List[Dim]]
        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 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_size_tensor(self) -> _t.Tensor:
        """
        :return: size tensor, or dyn_size_ext if defined
        :rtype: _t.Tensor
        """
        if self.dyn_size_ext:
            return self.dyn_size_ext

        import returnn.frontend as rf

        assert self.size is not None
        return rf.convert_to_tensor(self.size, name="%s:size" % self.description)

    def get_dim_value(self) -> Union[int, _t.RawTensorType]:
        """
        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).

        :return: max(size or dyn_size)
        """
        res = self.get_dim_value_tensor()
        if isinstance(res, _t.Tensor):
            assert res.dims == ()
            return res.raw_tensor
        assert isinstance(res, int)
        return res

    def get_dim_value_tensor(self) -> Union[int, _t.Tensor]:
        """
        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).

        :return: max(size or dyn_size)
        """
        import returnn.frontend as rf

        if self.dimension is not None:
            return self.dimension
        if self.dyn_size_ext and self.dyn_size_ext.placeholder is not None:  # fast path
            if self.dyn_size_ext.batch_ndim > 0:
                return rf.reduce_max(
                    self.dyn_size_ext,
                    axis=self.dyn_size_ext.dim_tags,
                    # Masking is not always possible here, e.g.
                    # self = Dim{'self-att-keys'['time:var:extern_data:classes'[B]]}.
                    use_mask=False,
                )
            return self.dyn_size_ext
        if self.is_batch_dim():
            res = None
            if self._extra and self._extra.src_data:
                res = self._extra.src_data.get_batch_dim()
            elif self.batch:
                res = self.batch.dim
            if isinstance(res, int):
                return res
            if res is not None:
                return _t.Tensor("batch", dims=(), dtype=rf.get_default_array_index_dtype(), raw_tensor=res)
        if (
            self._extra
            and self._extra.src_data is not None
            and self._extra.src_axis is not None
            and self._extra.src_data.placeholder is not None
        ):
            res = self._extra.src_data.get_dim(self._extra.src_axis)
            if isinstance(res, int):
                return res
            return _t.Tensor("batch", dims=(), dtype=rf.get_default_array_index_dtype(), raw_tensor=res)
        self.complete_dyn_size()
        if self.dyn_size_ext and self.dyn_size_ext.placeholder is not None:
            if self.dyn_size_ext.batch_ndim > 0:
                return rf.reduce_max(self.dyn_size_ext, axis=self.dyn_size_ext.dim_tags)
            return self.dyn_size_ext
        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]

    def _get_same_base_extra(self) -> Optional[_DimExtra]:
        if not self._extra:
            return None
        base = self.get_same_base()
        # noinspection PyProtectedMember
        return base._extra

    def _make_extra(self: _d.Dim) -> _DimExtra:
        if not self._extra:
            self._extra = _DimExtra(dim=self)
        return self._extra

    @property
    def vocab(self):
        """
        :rtype: returnn.datasets.util.vocabulary.Vocabulary|None
        """
        extra = self._get_same_base_extra()
        if extra:
            return extra.vocab
        return None

    @vocab.setter
    def vocab(self, vocab):
        """
        :param returnn.datasets.util.vocabulary.Vocabulary|None vocab:
        """
        if vocab is self.vocab:
            return
        if self.same_as:
            self.get_same_base().vocab = vocab
            return
        extra = self._get_same_base_extra()
        if extra:
            extra.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: Dim, other):
        """
        :param Dim|int other:
        :return: self + other. note that this is not commutative, i.e. different from other + self.
        :rtype: Dim
        """
        term = _OpLinearTerm.from_dim(self)
        term.extend_add_sub_(other, kind="add", right=True)
        return term.as_dim()

    def __radd__(self: Dim, other):
        """
        :param Dim|int other:
        :return: other + self
        :rtype: Dim
        """
        term = _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: Dim, other):
        """
        :param Dim|int other:
        :return: self - other
        :rtype: Dim
        """
        term = _OpLinearTerm.from_dim(self)
        term.extend_add_sub_(other, kind="sub", right=True)
        return term.as_dim()

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

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

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

    def __floordiv__(self: Dim, other):
        """
        :param Dim|int other:
        :rtype: Dim
        """
        term = _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: Dim, other):
        """
        :param Dim|int other:
        :rtype: Dim
        """
        term = _OpLinearTerm.from_dim(self)
        term.extend_mul_div_(other, kind="truediv", right=False)
        return term.as_dim()

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

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

    def ceildiv_right(self: Dim, other):
        """
        :param Dim|int other:
        :rtype: Dim
        """
        term = _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

    def is_constant_static_dim(self) -> bool:
        """
        :return: derived op of type constant
        """
        return self.derived_from_op and self.derived_from_op.kind == "constant"


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


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: Optional[_d.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, Op):
            return self._value() == other._value()
        return False

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


def _get_description(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 _OpMultTerm:
    """
    represents sth like a * b * c
    """

    @classmethod
    def from_dim(cls, dim: _d.Dim) -> _OpMultTerm:
        """
        :param dim:
        :return: op mult term
        """
        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, _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 _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 _OpMultTerm(self.terms[1:])
        res = self.copy()
        res.extend_mul_div_(_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
        if kind == "mul":
            # We want (b * a) // b != a.
            # However, we want h * (2 * a // h) == 2 * a.
            # So, for `mul`, and only for `mul`, check all terms, whether we can simplify some division-term.
            for i, term in reversed(list(enumerate(self.terms))) if right else enumerate(self.terms):
                assert isinstance(term, _d.Dim)
                if 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 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():
                    return i
        # For the last/first term, extra checks.
        i = len(self.terms) - 1 if right else 0
        term = self.terms[i]
        if kind.endswith("div") and other == term:
            return i
        op_kind = kind + "_" + ("right" if right else "left")
        if term.derived_from_op and term.derived_from_op.kind == op_kind:
            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 = [_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, _d.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] = _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] = _make_constant_static_dim(term.dimension // other.dimension, kind=term.kind)
                    return
                # Fallback with generic handling.
            op_kind = kind + "_" + ("right" if right else "left")
            if kind.endswith("div") and term.derived_from_op and term.derived_from_op.kind == op_kind:
                numerator = term.derived_from_op.inputs[0]
                denominator = term.derived_from_op.inputs[1]
                self.terms[idx] = _OpMultTerm.new_div_dim(numerator, denominator * other, kind=kind, right=right)
                return
        if kind.endswith("div"):
            self.terms = [_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" % (_get_description(numerator), _get_description(denominator))
        elif kind == "ceildiv" and right:
            description = "⌈%s/%s⌉" % (_get_description(numerator), _get_description(denominator))
        else:
            description = "%s_%s(%s, %s)" % (
                kind,
                "right" if right else "left",
                _get_description(numerator, brackets=False),
                _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 _d.Dim(
            description=description,
            kind=numerator.kind,
            dimension=dim_value,
            derived_from_op=Op(kind=op_kind, inputs=[numerator, denominator]),
            derived_from_tag=numerator,
        )

    def as_dim(self):
        """
        :rtype: Dim
        """
        if self.is_one():
            return _make_constant_static_dim(1)
        if len(self.terms) == 1:
            return self.terms[0]
        dim_kind = _get_merged_dim_kind(self.terms)
        return _d.Dim(
            kind=dim_kind,
            description="*".join(map(_get_description, self.terms)),
            dimension=self.dimension,
            derived_from_op=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 != DimTypes.Unspecified:
                return term_
        # Now find any.
        for term_ in self.terms:
            return term_
        return None


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 _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, _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 _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(_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 _d.Dim(
            kind=_get_merged_dim_kind(add_parts),
            description="+".join(desc_parts),
            dimension=dim,
            derived_from_op=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 = _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] = _OpMultTerm.from_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 = _OpMultTerm.from_dim_factors(most_recent_term.terms[:-1]).as_dim()
                b = _OpMultTerm.from_dim_factors(term.terms[:-1]).as_dim()
                res = _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 = _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") and len(self.terms) >= 2:
            if any(not term.divisible(other, right=right) for term in self.terms):
                self.terms = [
                    _OpMultTerm.from_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 _make_constant_static_dim(other, kind=base_tag.kind if base_tag else None)
        elif isinstance(other, _d.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 != DimTypes.Unspecified:
                    return term_
        # Now find any.
        for term in self.terms:
            for term_ in term.terms:
                return term_
        return None


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 DimTypes.Batch
    elif any(tag.is_feature_dim() for tag in dim_tags):
        return DimTypes.Feature
    else:
        return DimTypes.Spatial


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, _d.Dim):
        obj = obj.get_same_base()
        return (
            "",
            obj.description,
            obj.kind,
            obj.dimension,
            obj.dyn_size_ext.dims if obj.dyn_size_ext is not None else None,
        )
    if isinstance(obj, _m.MarkedDim):
        return obj.__class__.__name__, obj.tag
    return obj
