"""
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, Any, Tuple, Sequence, MutableMapping, Dict, List, Set, Callable
import operator
import weakref

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 returnn.datasets.util.vocabulary import Vocabulary

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", global_base=_d, global_name="Dim.Types.Batch")
    Spatial = Entity("spatial", global_base=_d, global_name="Dim.Types.Spatial")  # also time
    Time = Spatial  # we don't treat this as different
    Feature = Entity("feature", global_base=_d, global_name="Dim.Types.Feature")
    Types = (Batch, Spatial, Feature)


class _DimExtra:
    def __init__(
        self,
        *,
        dim: Dim,
        kind: Entity = DimTypes.Unspecified,
        vocab: Union[None, Dict[str, Any], Vocabulary] = None,
        undefined: bool = False,
        special: bool = False,
        auto_generated: bool = False,
        match_priority: int = 0,
        derived_from_tag: Optional[Dim] = None,
        derived_from_op: Optional[Op] = None,
        batch: Optional[BatchInfo] = None,
        control_flow_ctx: Optional[ControlFlowContext] = None,
        src_data: Optional[_t.Tensor] = None,
        src_axis: Optional[int] = None,
    ):
        """
        :param dim:
        :param kind:
        :param vocab:
        :param undefined: When this is specified as `None` by the user via `shape`.
        :param 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 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 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 derived_from_op:
        :param 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 batch: for batch-dim, or dynamic dims per batch
        :param 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
        if vocab:
            from returnn.datasets.util.vocabulary import Vocabulary

            if isinstance(vocab, Vocabulary):
                pass
            elif isinstance(vocab, dict):
                vocab = Vocabulary.create_vocab(**vocab)
            else:
                raise TypeError(f"invalid vocab {vocab!r} type {type(vocab)}")
        self.vocab = vocab
        self.same_as: Optional[Dim] = None
        self.copy_same_as: Optional[Dim] = None
        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 = dim
        self.match_priority = match_priority
        if src_data is not None:
            assert isinstance(src_data, _t.Tensor) and isinstance(src_axis, int)
        if not batch and dim.dyn_size_ext is not None:
            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 RefIdEq (earlier 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: Dict[Tuple[BatchInfo, Optional[ControlFlowContext]], Dim] = {}
        self.cache_dyn_size_ext_dev: Dict[str, _t.Tensor] = {}  # device -> dyn_size_ext
        self.cache_seq_mask: Dict[Tuple[str, Optional[Tuple[Dim, ...]]], _t.Tensor] = {}  # (dev,dim_order) -> seq_mask
        self.cache_dim_math = _CacheDimMath()  # op (add,sub,...), operand -> Dim

    def __getstate__(self):
        d = vars(self).copy()
        d["batch"] = None
        d["same_for_batch_ctx"] = {}
        d["cache_dyn_size_ext_dev"] = {}
        d["cache_seq_mask"] = {}
        d["cache_dim_math"] = {}
        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:
            # noinspection PyTypeChecker
            self.kind = {v.name: v for v in DimTypes.Types}[self.kind]

    def __sis_state__(self):
        raise ValueError(f"{self}: currently not expected to be part of the Sisyphus state/hash")


class _DimMixin:
    name: Optional[str]
    capacity: Optional[int]
    size: Optional[int]
    dyn_size_ext: Optional[_t.Tensor]
    _dyn_size_max_value: Optional[_t.Tensor]  # scalar
    _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 self._extra.vocab and self.size is None:
                self.size = self._extra.vocab.num_labels
                self.capacity = self.capacity or self.size
        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

    # This is potentially replaced by native implementation (Dim.size.__get__).
    @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 is not None:
                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 is not None:
                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 _d.batch_dim:
            return "B"  # Data.__repr__ will additionally give info on the batch
        desc_ = self.get_same_base().description
        desc = {DimTypes.Feature: "F", DimTypes.Batch: "B"}.get(self.kind, "")
        desc += repr(desc_) if desc_ is not None else "-"
        if self.special:
            desc += "!"
        elif self.dimension is not None:
            desc += f"({self.dimension})"
        else:
            if self.dyn_size_ext is not None:
                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"
        func, args, (vs, slots), *more_args = super().__reduce_ex__(protocol)
        assert not vs
        assert isinstance(slots, dict) and "_dyn_size_max_value" in slots
        slots["_dyn_size_max_value"] = None
        # noinspection PyRedundantParentheses
        return (func, args, (vs, slots), *more_args)

    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 reset_eager(self: Dim):
        """
        In an eager-based framework, dyn_size_ext.raw_tensor etc will be different in each step.
        This resets everything related.
        This can also include caches.
        """
        self.reset_raw()

    def reset_raw(self: Dim, *, only_self: bool = False, include_parents: bool = False):
        """
        Reset all raw tensors.
        """
        visited = set()  # ids
        queue = [self]
        while queue:
            dim: Dim = queue.pop()
            if id(dim) in visited:
                continue
            visited.add(id(dim))
            dim.reset_batch_ctx()
            dim._dyn_size_max_value = None
            if dim.dyn_size_ext is not None:
                dim.dyn_size_ext.reset()
            # noinspection PyProtectedMember
            dim_extra = dim._extra
            if dim_extra:
                dim_extra.cache_dyn_size_ext_dev.clear()
                dim_extra.cache_seq_mask.clear()
                if dim.dyn_size_ext is not None or dim.dimension is None:
                    dim_extra.cache_dim_math.clear()
                else:
                    dim_extra.cache_dim_math.clear_dynamic()
            if only_self:
                return
            if dim_extra:
                # Any dims via dim math could also contain raw tensors,
                # so iterate through them.
                queue += dim_extra.cache_dim_math.values()
                if dim_extra.same_as:
                    queue.append(dim_extra.same_as)
                if dim_extra.copy_same_as:
                    queue.append(dim_extra.copy_same_as)
                queue += dim_extra.same_for_batch_ctx.values()
                if include_parents and dim_extra.derived_from_op:
                    queue.extend(dim_extra.derived_from_op.inputs)

    def reset_batch_and_raw(self: Dim):
        """
        Reset batch and raw tensors.
        """
        self.reset_raw(include_parents=True)

    def transform_tensors(self: Dim, func: Callable[[_t.Tensor], None]):
        """
        Transforms all tensors inplace, e.g. Numpy to PyTorch or so.
        Resets all caches.

        :param func: operates inplace
        """
        dyn_size_ext = self.dyn_size_ext.copy() if self.dyn_size_ext is not None else None
        self.reset_raw(only_self=True)
        if dyn_size_ext is not None:
            func(dyn_size_ext)
        self.dyn_size_ext = dyn_size_ext

    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 self.dyn_size_ext is None:
            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: Dim):
        """
        :rtype: bool
        """
        if (self.dyn_size_ext is not None and not self.dyn_size_ext.is_valid_in_current_graph()) or (
            self._dyn_size_max_value is not None and not self._dyn_size_max_value.is_valid_in_current_graph()
        ):  # maybe from an earlier run which reuses the dim tag
            # Reset and cleanup.
            self.reset_batch_ctx()
            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 is not None 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 is not None 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 is not None and self.dyn_size_ext is None:
                    self.dyn_size_ext = same.dyn_size_ext
                if same.dyn_size_ext is not None 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 is not None and same.dyn_size_ext is None:
                    same.dyn_size_ext = self.dyn_size_ext
                if self.dyn_size_ext is not None and self.dyn_size_ext.placeholder is not None:
                    if same.dyn_size_ext is None or same.dyn_size_ext.placeholder is None:
                        same.dyn_size_ext = self.dyn_size_ext
                # noinspection PyProtectedMember
                if self._dyn_size_max_value is None and same._dyn_size_max_value is not None:
                    # noinspection PyProtectedMember
                    self._dyn_size_max_value = same._dyn_size_max_value
                # noinspection PyProtectedMember
                if same._dyn_size_max_value is None and self._dyn_size_max_value is not None:
                    # noinspection PyProtectedMember
                    same._dyn_size_max_value = self._dyn_size_max_value

    def get_for_batch_ctx(
        self: Dim, batch: BatchInfo, ctx: Optional[ControlFlowContext], *, allow_none: bool = 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:
        """
        from returnn.tensor import ControlFlowContext

        assert self.can_be_used_as_dim()
        if self.batch == batch and self._can_use_in_ctx(ctx) and self.dyn_size_ext is not None:
            self._validate_in_current_graph()
            self._maybe_update()
            if self.batch == batch and self._can_use_in_ctx(ctx) and self.dyn_size_ext is not None:  # 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 batch.batch_dim_tag
        if self.is_static():
            # If static dim, no effect.
            assert not self.batch
            return self
        if batch.is_broadcast():
            return self  # just leave as-is. should not matter.
        dim_tag = None
        if self._extra:
            same_base = self.get_same_base()
            same_base._validate_in_current_graph()
            # noinspection PyProtectedMember
            if same_base._extra:
                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
                        if tag.dyn_size_ext is not None:
                            return tag
                        dim_tag = tag
                        break
            if same_base.batch == batch and same_base._can_use_in_ctx(ctx) and same_base.dyn_size_ext is not None:
                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,))
        if dim_tag:
            assert dim_tag.dyn_size_ext is None
        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 is not None:
                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 is not None
                    ):
                        base_can_use_in_ctx = tag
                        break
            if base_can_use_in_ctx and base_can_use_in_ctx.dyn_size_ext is not None:
                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 is not None
                        # 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 dyn_size_ext is None:
            # Maybe we can infer dyn_size_ext, even with different batch.
            # Keep logic in sync with is_dim_known_in_batch_ctx.
            candidates = [self, same_base] + list(same_base_extra.same_for_batch_ctx.values())
            for other in candidates:
                if other.dyn_size_ext is not None and ControlFlowContext.is_parent_or_same(other.control_flow_ctx, ctx):
                    dyn_size_ext = other.dyn_size_ext.copy_template()
                    dyn_size_ext.beam = batch.beam
                    dyn_size_ext.batch = batch
                    break
        if dyn_size_ext is not None:
            ctx = dyn_size_ext.control_flow_ctx
        elif dim_tag:
            ctx = dim_tag.control_flow_ctx
        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 is not None:
                    if candidate.dyn_size_ext is not None:
                        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 is not None:
                    candidate.dyn_size_ext.batch = batch
                else:
                    candidate.complete_dyn_size(template_only=True)
                if not dim_tag:
                    dim_tag = candidate
        if not dim_tag:
            if allow_none:
                return None
            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 is not None 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
        if dyn_size_ext is not None:
            if dim_tag.dyn_size_ext is None:
                dim_tag.dyn_size_ext = dyn_size_ext
            else:
                assert dim_tag.dyn_size_ext.dims == dyn_size_ext.dims
        elif dim_tag.dyn_size_ext is not None:
            pass
        else:
            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._extra.cache_seq_mask.clear()
            self._extra.cache_dyn_size_ext_dev.clear()
        self.batch = None
        self.control_flow_ctx = None
        if self.dyn_size_ext is not None and self.dyn_size_ext.batch:
            self.dyn_size_ext = self.dyn_size_ext.copy_template()
            self.dyn_size_ext.batch = None
            self.dyn_size_ext.control_flow_ctx = None
        self._dyn_size_max_value = 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 is not None:
            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: _t.Tensor|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 is not None:
            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 is not None 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
            assert self.dyn_size_ext.placeholder is dyn_size
            return
        self._init_default_dyn_size_ext(dyn_size)
        self.set_tag_on_size_tensor(dyn_size)
        assert self.dyn_size_ext is not None and self.dyn_size_ext.placeholder is dyn_size

    def _init_default_dyn_size_ext(self, dyn_size):
        """
        :param tf.Tensor dyn_size:
        """
        if self.dyn_size_ext is not None:
            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_dyn_size_ext_for_device(self: Dim, device: Optional[str]) -> _t.Tensor:
        """
        :return: dyn_size_ext on the device
        """
        assert self.dyn_size_ext is not None
        if not device or device == "cpu":
            return self.dyn_size_ext

        import returnn.frontend as rf

        self._make_extra()
        if device in self._extra.cache_dyn_size_ext_dev:
            return self._extra.cache_dyn_size_ext_dev[device]
        self._extra.cache_dyn_size_ext_dev[device] = rf.copy_to_device(self.dyn_size_ext, device=device)
        return self._extra.cache_dyn_size_ext_dev[device]

    def get_mask(self: Dim, *, dim_order: Optional[Sequence[Dim]] = None, device: Optional[str] = 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.
        :param str|None device: if given, will move the mask to this device
        :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 is not None and self.dyn_size_ext.raw_tensor is not None
        # noinspection PyProtectedMember
        backend = self.dyn_size_ext._raw_backend

        if not device:
            device = rf.get_default_device()

        self._make_extra()
        dim_order_default = self.dyn_size_ext.dims + (self,)
        if dim_order is not None:
            dim_order = tuple([d for d in dim_order if d in dim_order_default])  # filter
        else:
            dim_order = dim_order_default
        cache_key = (device, dim_order)
        if cache_key in self._extra.cache_seq_mask:
            return self._extra.cache_seq_mask[cache_key]

        if self._extra.copy_same_as:
            if dim_order:
                dim_order = tuple(self._extra.copy_same_as if d == self else d for d in dim_order)
            mask = self._extra.copy_same_as.get_mask(dim_order=dim_order, device=device)
            mask, _ = rf.replace_dim(mask, in_dim=self._extra.copy_same_as, out_dim=self)
            return mask

        size_ext = self.get_dyn_size_ext_for_device(device)

        max_idx = rf.reduce(
            size_ext,
            axis=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 max_idx.
        # This assumes that max_idx >= 1.
        size_ext = size_ext.copy_masked(max_idx)
        idx_range = backend.range_over_dim(self, device=device)
        seq_mask = rf.compare(idx_range, "<", size_ext, allow_broadcast_all_sources=True, dim_order=dim_order)
        self._extra.cache_seq_mask[cache_key] = seq_mask
        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 self.dyn_size_ext is None and self.dimension is not None:
            return True
        if self.dyn_size_ext is not None:
            return True
        extra = self._get_same_base_extra()
        if extra:
            for _, other in extra.same_for_batch_ctx.items():
                if other.dyn_size_ext is not None:
                    return True
        return False

    def is_dim_known_in_batch_ctx(self: Dim, batch: BatchInfo, ctx: Optional[ControlFlowContext]) -> bool:
        """
        :return: whether :func:`get_for_batch_ctx` would return a valid existing dim tag
        """
        from returnn.tensor import ControlFlowContext

        if self.is_batch_dim():
            return True
        if self.is_static():
            return True
        dim = self.get_for_batch_ctx(batch=batch, ctx=ctx, allow_none=True)
        if dim:
            return dim.dyn_size_ext is not None
        candidates = [self, self.get_same_base()]
        if self._extra:
            candidates += list(self._extra.same_for_batch_ctx.values())
        for dim in candidates:
            # By intention, ignore the batch, only check the ctx.
            # Keep logic in sync with get_for_batch_ctx.
            if ControlFlowContext.is_parent_or_same(dim.control_flow_ctx, ctx) and dim.dyn_size_ext is not None:
                return True
        return False

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

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

    def is_static(self) -> bool:
        """
        :return: static
        """
        return self.dimension is not None

    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 self.dyn_size_ext is None:  # unknown, so we can only guess
            if self.is_batch_dim():
                return False
            return True
        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 is not None 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 is not None 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, _backend=None):
        """
        In case we can calculate the dyn size, do that now.

        :param bool template_only:
        :param _backend:
        """
        if self.is_static():
            return
        self._validate_in_current_graph()
        if self.dyn_size_ext is not None 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_dim in op.inputs:
            if self.batch:
                x_dim = x_dim.get_for_batch_ctx(self.batch, self.control_flow_ctx)
            x_dim.complete_dyn_size(template_only=template_only)

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

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

        import numpy
        import returnn.frontend as rf
        import contextlib

        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")]
        y_name = self.description + ":seq-length"

        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

        if tf:
            _ctx_for_inputs = tf_util.same_control_flow_ctx

        else:

            @contextlib.contextmanager
            def _ctx_for_inputs(_arg):
                yield

        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))
            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
                if util.is_onnx_export_global():
                    return tf_util.onnx_compat_floor_div(a, b)
                return a // b
            elif kind == "ceildiv":
                if util.is_onnx_export_global():
                    return -tf_util.onnx_compat_floor_div(-a, b)
                return -(-a // b)
            else:
                raise ValueError("unknown op kind %r" % op.kind)

        def _bin_op(a, b):
            if a is None:
                if isinstance(b, int):
                    if not template_only and backend and not tf:
                        return backend.convert_to_tensor(b, dims=(), dtype=size_dtype, name=y_name, device="cpu")
                    else:
                        y__ = _t.Tensor(name=y_name, dims=(), dtype=size_dtype)
                        if not template_only and tf:
                            with tf.control_dependencies(None):  # this will reset the context
                                y__.raw_tensor = tf.constant(b)
                        return y__
                elif isinstance(b, _t.Tensor):
                    return b.copy(name=y_name)
                else:
                    raise TypeError(f"complete_dyn_size: _bin_op: unexpected type {type(b)}")
            assert isinstance(a, _t.Tensor)
            if template_only or not backend:
                if isinstance(b, _t.Tensor):
                    return _t.Tensor.get_common_data([a, b], allow_broadcast_all_sources=True)
                assert isinstance(b, int)
                return a.copy_template()
            if tf:
                if isinstance(b, _t.Tensor):
                    res = _t.Tensor.get_common_data([a, b], allow_broadcast_all_sources=True)
                    with _ctx_for_inputs(a):
                        a = a.copy_compatible_to_dims(res.dims) if a.dims else a
                    with _ctx_for_inputs(b):
                        b = b.copy_compatible_to_dims(res.dims) if b.dims else b
                else:
                    assert isinstance(b, int)
                    res = a.copy_template()
                with _ctx_for_inputs([a, b]):
                    res.raw_tensor = _bin_op_tf(a.raw_tensor, b.raw_tensor if isinstance(b, _t.Tensor) else b)
                return res
            if kind == "add":
                return _relu(rf.combine_bc(a, "add", b))
            elif kind == "sub":
                return _relu(rf.combine_bc(a, "sub", b))
            elif kind == "mul":
                return rf.combine_bc(a, "mul", b)
            elif kind in ("floordiv", "truediv"):  # truediv assumes there is no remainder
                return rf.combine_bc(a, "floordiv", b)
            elif kind == "ceildiv":
                return -rf.combine_bc(-a, "floordiv", b)
            else:
                raise ValueError("unknown op kind %r" % op.kind)

        def _relu(a):
            if isinstance(a, _t.Tensor):
                return rf.relu(a)
            elif isinstance(a, int):
                return max(a, 0)
            else:
                raise TypeError(f"complete_dyn_size: _relu: unexpected type {type(a)}")

        y: Optional[_t.Tensor] = None  # resulting dyn size
        inputs = list(op.inputs)
        assert inputs
        for x_dim in inputs:
            x_dim: Dim
            if self.batch:
                x_dim = x_dim.get_for_batch_ctx(self.batch, self.control_flow_ctx)
            x_dim.complete_dyn_size(template_only=template_only, _backend=backend)
            if x_dim.dyn_size_ext is None and x_dim.dimension is None:
                return
            y = _bin_op(y, x_dim.dimension if x_dim.dimension is not None else x_dim.dyn_size_ext)
        assert y is not None, f"op {op}?"
        if self.dyn_size_ext is not None:
            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 not template_only and y.raw_tensor is not None:
            # Note: Earlier, we had this wrong.
            # It is not correct to replicate the same math (bin ops)
            # on the dim values (_dyn_size_max_value of each dim).
            # Consider sizes1=[2,3], sizes2=[5,4], and the op is "add".
            # Then the result sizes would be [7,7], thus its max is 7,
            # but max(sizes1)+max(sizes2)=3+5=8.
            self._dyn_size_max_value = rf.reduce_max(y, axis=y.dims) if y.dims else y
        if tf and y.placeholder is not None:
            self.set_tag_on_size_tensor(y.placeholder)

    # Set by more recent behavior versions.
    _SimpleEquality = False

    def is_equal(
        self: Dim,
        other: Dim,
        *,
        ignore_feature_dim: bool = False,
        allow_same_feature_dim: bool = False,
        allow_same_spatial_dim: Optional[bool] = None,
        treat_feature_as_spatial: bool = False,
        broadcast_matches: bool = False,
        unknown_spatial_matches: bool = False,
        undefined_matches: bool = False,
        derived_matches: bool = False,
        allow_old_behavior: bool = False,
    ) -> bool:
        """
        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 other:
        :param ignore_feature_dim:
        :param allow_same_feature_dim:
        :param allow_same_spatial_dim:
        :param treat_feature_as_spatial:
        :param broadcast_matches:
        :param unknown_spatial_matches:
        :param undefined_matches:
        :param derived_matches:
        :param allow_old_behavior: useful e.g. for find_matching_dim_map
        """
        if self is other:  # first some fast path check
            return True
        if not isinstance(other, _d.Dim):
            return False
        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(same_dim=True) if derived_matches else self.get_same_base()
        other_base = other.get_same_derived_base(same_dim=True) if derived_matches else other.get_same_base()
        if self_base is other_base:
            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:
            # Note: No kind at all: Reinterpret treat_feature_as_spatial a bit:
            # Assume that we want them all to be handled the same, no matter the kind.
            # (Except of batch dim kind, which is still excluded here.)
            if self_kind == DimTypes.Feature or not self_kind:
                self_kind = DimTypes.Spatial
            if other_kind == DimTypes.Feature or not other_kind:
                other_kind = DimTypes.Spatial
        if self.dimension != other.dimension:
            if broadcast_matches and (
                # Only auto-generated dim tags are allowed to be treated as broadcastable.
                # This was another suggestion from here: https://github.com/rwth-i6/returnn/issues/666
                # It was not implemented like this because the auto_generated flag was only introduced later.
                (self.dimension == 1 and self.auto_generated) or (other.dimension == 1 and other.auto_generated)
            ):
                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 self._SimpleEquality and not allow_old_behavior:
            # 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: Dim, other: Dim) -> bool:
        """
        :param other:
        :return: :func:`is_equal` with default options
        """
        if self is other:  # fast path
            return True
        if not isinstance(other, _d.Dim):
            return False
        if self._SimpleEquality:  # fast path
            return self._eq_simple(other)
        return self.is_equal(other)

    def _eq_simple(self: Dim, other: Dim) -> bool:
        if self is other:  # fast path
            return True
        if not isinstance(other, _d.Dim):
            return False
        # See is_equal for the logic. This here should exactly replicate it.
        # Inline self = self.get_same_base().
        # noinspection PyProtectedMember
        while self._extra and self._extra.same_as:
            # noinspection PyProtectedMember,PyMethodFirstArgAssignment
            self = self._extra.same_as
        # Inline other = other.get_same_base().
        # noinspection PyProtectedMember
        while other._extra and other._extra.same_as:
            # noinspection PyProtectedMember
            other = other._extra.same_as
        if self is other:
            return True
        # noinspection PyProtectedMember
        if self._extra and other._extra:
            self_extra = self._extra
            # noinspection PyProtectedMember
            other_extra = other._extra
            if self_extra.kind == other_extra.kind == DimTypes.Batch:
                return True
            # noinspection PyProtectedMember
            if self_extra.auto_generated and other_extra.auto_generated and self.description == other.description:
                return True
        return False

    def __ne__(self: Dim, other: Dim) -> bool:
        """
        :param other:
        """
        if self is other:  # fast path
            return False
        if not isinstance(other, _d.Dim):
            return True
        if self._SimpleEquality:  # fast path
            return not self._eq_simple(other)
        return not self.is_equal(other)

    def _ne_simple(self: Dim, other: Dim) -> bool:
        return not self._eq_simple(other)

    def _ne_generic(self: Dim, other: Dim) -> bool:
        return not self.is_equal(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).
        # Inline self = self.get_same_base().
        self_extra = self._extra
        # noinspection PyProtectedMember
        while self_extra and self_extra.same_as:
            # noinspection PyMethodFirstArgAssignment
            self = self_extra.same_as
            self_extra = self._extra
        if self_extra:
            if self_extra.special:
                return hash(id(self))
            if self_extra.kind == DimTypes.Batch:
                return hash(())
            if self_extra.auto_generated:
                return hash((self_extra.kind, self.size, self.name))
        return hash(id(self))

    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, *, same_dim: bool = False) -> _d.Dim:
        """
        :param same_dim: if True, return the last base which has the same dimension.
            The derived base might have a different dimension.
            In case it is dynamic, the dimension is None, so it is always the same.
            In case it is static, there might be a different dimension.
        :return: same base, but also consider derived_from_...
        """
        last_base_self_dim = self
        base = self
        visited = {}
        while base.same_as or base.derived_from_tag:
            assert id(base) not in visited  # should not have cycles. normally this should never be triggered
            visited[id(base)] = base
            if base.same_as:
                base = base.same_as
                continue
            base = base.derived_from_tag
            assert base
            if base.dimension == self.dimension:
                last_base_self_dim = base
        return last_base_self_dim if same_dim else 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._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 []):
                dim_.reset_raw()
        other_same_base._merge_same_for_batch_ctx_dict(self)
        # noinspection PyProtectedMember
        if self_same_as._extra:
            # noinspection PyProtectedMember
            for k, v in self_same_as._extra.cache_dim_math.items():
                # noinspection PyProtectedMember
                other_same_base._extra.cache_dim_math.setdefault(k, v)
            # noinspection PyProtectedMember
            self_same_as._extra.cache_dim_math.clear()
        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 (
                tag.dyn_size_ext is None 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 self.dyn_size_ext is None 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.dyn_size_ext is not None
            and self.dyn_size_ext.raw_tensor is None
            and other_same_base.dyn_size_ext.raw_tensor is not None
        ):
            self.dyn_size_ext = other_same_base.dyn_size_ext.copy()
            self._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 other_same_base.is_static():
            # 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
        if self_same_as is not self:
            assert not self_same_as.same_as
            assert self_same_as is not other_same_base
            with util.guard_infinite_recursion(_d.Dim.declare_same_as, self_same_as, other_same_base):
                # Do this at the end, when this is not our own base anymore,
                # because otherwise any reset() might incorrectly reset the new other base.
                self_same_as.declare_same_as(other_same_base)

    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 is not None or dim.dyn_size_ext is None):
                    continue  # keep ours
                if dim.dyn_size_ext is None:
                    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 self.is_dynamic() or not self.is_dim_known():
            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.dyn_size_ext is None:
                if base.dyn_size_ext is not None:
                    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)
        self._make_extra().copy_same_as = other

    @classmethod
    def get_existing_tag_from_collection(
        cls, other: Dim, tags: Union[Sequence[Dim], Set[Dim]], is_equal_opts: Optional[Dict[str, Any]] = None
    ) -> Optional[Dim]:
        """
        :param other:
        :param tags:
        :param is_equal_opts: passed to Dim.is_equal
        """
        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: List[_t.Tensor],
        is_equal_opts: Optional[Dict[str, Any]] = None,
        unique_separate_axes: bool = True,
    ) -> Tuple[List[Dim], util.DictRefKeys[_t.Tensor, List[Dim]]]:
        """
        :param data_list:
        :param is_equal_opts: passed to Dim.is_equal
        :param 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)
        """
        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, *, device: Optional[str] = None) -> _t.Tensor:
        """
        :param device: if None, will use CPU
        :return: size tensor, or dyn_size_ext if defined
        :rtype: _t.Tensor
        """
        if self.dyn_size_ext is not None:
            if not device or device == "cpu":
                return self.dyn_size_ext
            return self.get_dyn_size_ext_for_device(device)

        import returnn.frontend as rf

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

    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 == ()
            assert res.raw_tensor is not None
            return res.raw_tensor
        assert isinstance(res, int)
        return res

    def get_dim_value_tensor(self: Dim) -> 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_max_value is not None:  # fast path, precomputed
            assert self._dyn_size_max_value.raw_tensor is not None
            return self._dyn_size_max_value
        if (
            self._extra
            and self._extra.src_data
            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
            res = _t.Tensor(
                f"{self._extra.src_data}:shape[{self._extra.src_axis}]",
                dims=(),
                dtype=rf.get_default_array_index_dtype(),
                raw_tensor=res,
            )
            self._dyn_size_max_value = res
            return res
        self.complete_dyn_size()
        if self._dyn_size_max_value is not None:
            return self._dyn_size_max_value
        if self.dyn_size_ext is not None and self.dyn_size_ext.placeholder is not None:
            if self.dyn_size_ext.batch_ndim > 0:
                res = 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,
                )
            else:
                res = self.dyn_size_ext.copy()
            assert res.raw_tensor is not None
            self._dyn_size_max_value = res
            return res
        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)
        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
        if vocab:
            self.get_same_base()._make_extra().vocab = vocab
        else:
            self._get_same_base_extra().vocab = None

    # 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
        """
        if _is_const_dim_value(other, 0):
            return self
        cache_key = ("add", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = _MathFindMatchingAdditive(start=self, right=True, other=other).search_and_maybe_replace()
        if not res:
            res = _math_get_dim_via_bin_op([self, other], "add")
        cache[cache_key] = res
        return res

    def __radd__(self: Dim, other):
        """
        :param Dim|int other:
        :return: other + self
        :rtype: Dim
        """
        if _is_const_dim_value(other, 0):
            return self
        cache_key = ("add_left", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = _MathFindMatchingAdditive(start=self, right=False, other=other).search_and_maybe_replace()
        if not res:
            res = _math_get_dim_via_bin_op([other, self], "add")
        cache[cache_key] = res
        return res

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

    def sub_right(self: Dim, other):
        """
        :param Dim|int other:
        :return: self - other
        :rtype: Dim
        """
        if _is_const_dim_value(other, 0):
            return self
        if (
            self.derived_from_op
            and self.derived_from_op.kind == "add"
            and len(self.derived_from_op.inputs) == 2
            and _dim_or_const_equal(other, self.derived_from_op.inputs[1])
        ):
            return self.derived_from_op.inputs[0]
        cache_key = ("sub", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = self + (-other)
        cache[cache_key] = res
        return res

    def sub_left(self: Dim, other):
        """
        :param Dim|int other:
        :return: (-other) + self
        :rtype: Dim
        """
        if _is_const_dim_value(other, 0):
            return self
        if (
            self.derived_from_op
            and self.derived_from_op.kind == "add"
            and len(self.derived_from_op.inputs) == 2
            and _dim_or_const_equal(other, self.derived_from_op.inputs[0])
        ):
            return self.derived_from_op.inputs[1]
        cache_key = ("sub_left", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = (-other) + self
        cache[cache_key] = res
        return res

    def __mul__(self: Dim, other):
        """
        :param Dim|int other:
        :rtype: Dim
        """
        if isinstance(other, _d.Dim) and other.is_constant_static_dim():
            other = other.dimension  # makes matching easier
        if isinstance(other, int) and other == 1:
            return self
        if self.is_constant_static_dim() and isinstance(other, _d.Dim) and not other.is_constant_static_dim():
            return self.dimension * other  # use rmul
        cache_key = ("mul", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = _math_find_matching_mult(start=self, right=True, other=other)
        if not res:
            res = _math_get_dim_via_bin_op([self, other], "mul")
        cache[cache_key] = res
        return res

    def __rmul__(self: Dim, other):
        """
        :param Dim|int other:
        :rtype: Dim
        """
        if isinstance(other, _d.Dim) and other.is_constant_static_dim():
            other = other.dimension  # makes matching easier
        if isinstance(other, int) and other == 1:
            return self
        cache_key = ("mul_left", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = _math_find_matching_mult(start=self, right=False, other=other)
        if not res:
            res = _math_get_dim_via_bin_op([other, self], "mul")
        cache[cache_key] = res
        return res

    def __floordiv__(self: Dim, other):
        """
        :param Dim|int other:
        :rtype: Dim
        """
        if isinstance(other, _d.Dim) and other.is_constant_static_dim():
            other = other.dimension  # makes matching easier
        if isinstance(other, int) and other == 1:
            return self
        if (
            self.derived_from_op
            and self.derived_from_op.kind == "mul"
            and len(self.derived_from_op.inputs) == 2
            and _dim_or_const_equal(other, self.derived_from_op.inputs[1])
        ):
            return self.derived_from_op.inputs[0]
        cache_key = ("floordiv", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = _math_find_matching_div(start=self, right=True, other=other, kind="floordiv")
        if not res:
            res = _math_get_dim_via_bin_op([self, other], "floordiv")
        cache[cache_key] = res
        return res

    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
        """
        if isinstance(other, _d.Dim) and other.is_constant_static_dim():
            other = other.dimension  # makes matching easier
        if isinstance(other, int) and other == 1:
            return self
        if (
            self.derived_from_op
            and self.derived_from_op.kind == "mul"
            and len(self.derived_from_op.inputs) == 2
            and _dim_or_const_equal(other, self.derived_from_op.inputs[0])
        ):
            return self.derived_from_op.inputs[1]
        cache_key = ("truediv_left", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = _math_find_matching_div(start=self, right=True, other=other, kind="truediv_left")
        if not res:
            res = _math_get_dim_via_bin_op([self, other], "truediv_left")
        cache[cache_key] = res
        return res

    def div_right(self: Dim, other):
        """
        :param Dim|int other:
        :rtype: Dim
        """
        if isinstance(other, _d.Dim) and other.is_constant_static_dim():
            other = other.dimension  # makes matching easier
        if isinstance(other, int) and other == 1:
            return self
        if (
            self.derived_from_op
            and self.derived_from_op.kind == "mul"
            and len(self.derived_from_op.inputs) == 2
            and _dim_or_const_equal(other, self.derived_from_op.inputs[1])
        ):
            return self.derived_from_op.inputs[0]
        cache_key = ("truediv", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = _math_find_matching_div(start=self, right=True, other=other, kind="truediv")
        if not res:
            res = _math_get_dim_via_bin_op([self, other], "truediv")
        cache[cache_key] = res
        return res

    def ceildiv_left(self: Dim, other):
        """
        :param Dim|int other:
        :rtype: Dim
        """
        if isinstance(other, _d.Dim) and other.is_constant_static_dim():
            other = other.dimension  # makes matching easier
        if isinstance(other, int) and other == 1:
            return self
        if (
            self.derived_from_op
            and self.derived_from_op.kind == "mul"
            and len(self.derived_from_op.inputs) == 2
            and _dim_or_const_equal(other, self.derived_from_op.inputs[0])
        ):
            return self.derived_from_op.inputs[1]
        cache_key = ("ceildiv_left", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = _math_find_matching_div(start=self, right=True, other=other, kind="ceildiv_left")
        if not res:
            res = _math_get_dim_via_bin_op([self, other], "ceildiv_left")
        cache[cache_key] = res
        return res

    def ceildiv_right(self: Dim, other):
        """
        :param Dim|int other:
        :rtype: Dim
        """
        if isinstance(other, _d.Dim) and other.is_constant_static_dim():
            other = other.dimension  # makes matching easier
        if isinstance(other, int) and other == 1:
            return self
        if (
            self.derived_from_op
            and self.derived_from_op.kind == "mul"
            and len(self.derived_from_op.inputs) == 2
            and _dim_or_const_equal(other, self.derived_from_op.inputs[1])
        ):
            return self.derived_from_op.inputs[0]
        cache_key = ("ceildiv", other)
        cache = self.get_same_base()._make_extra().cache_dim_math
        cache_entry = cache.get(cache_key, None)
        if cache_entry:
            cache_entry.complete_dyn_size()
            return cache_entry
        res = _math_find_matching_div(start=self, right=True, other=other, kind="ceildiv")
        if not res:
            res = _math_get_dim_via_bin_op([self, other], "ceildiv")
        cache[cache_key] = res
        return res

    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 _cache_dim_math_get(
        self: Dim, op_kind: str, operand: Union[Dim, int]
    ) -> Tuple[_CacheDimMath, Tuple[str, Union[Dim, int]], Optional[Dim]]:
        same_base = self.get_same_base()
        # noinspection PyProtectedMember
        extra = same_base._make_extra()
        if isinstance(operand, _d.Dim) and operand.is_constant_static_dim():
            operand = operand.dimension
        cache = extra.cache_dim_math
        cache_key = (op_kind, operand)
        return cache, cache_key, cache.get(cache_key, None)

    def __sis_state__(self):
        # Note: The name (description) is deliberately not part of the state/hash,
        # as this is quite arbitrary set by the user,
        # and could easily change.
        # In general, we keep the relevant state for Sisyphus quite minimal.
        value = {"dim": self.dimension}
        if self.kind is not None:
            value["kind"] = self.kind.name
        return value


def _make_constant_static_dim(value, kind=None):
    """
    :param int value:
    :param Entity|None kind:
    :rtype: Dim
    """
    assert isinstance(value, int)
    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,
    )


def _dim_or_const_equal(a: Union[Dim, int], b: Union[Dim, int]) -> bool:
    if isinstance(a, int):
        if isinstance(b, int):
            return a == b
        elif isinstance(b, _d.Dim):
            return a == b.dimension and b.is_constant_static_dim()
        else:
            raise TypeError(f"unexpected b type {type(b)}")
    elif isinstance(a, _d.Dim):
        if a.is_constant_static_dim():
            if isinstance(b, int):
                return a.dimension == b
            elif isinstance(b, _d.Dim):
                return a.dimension == b.dimension and b.is_constant_static_dim()
            else:
                raise TypeError(f"unexpected b type {type(b)}")
        else:
            if isinstance(b, int):
                return False
            elif isinstance(b, _d.Dim):
                return a == b
            else:
                raise TypeError(f"unexpected b type {type(b)}")
    else:
        raise TypeError(f"unexpected a type {type(a)}")


_BinOps = {
    "add": operator.add,
    "mul": operator.mul,
    "sub": operator.sub,
    "floordiv": operator.floordiv,
    "truediv": operator.floordiv,
    "truediv_left": operator.floordiv,
    "ceildiv": lambda a, b: -(-a // b),
    "ceildiv_left": lambda a, b: -(-a // b),
}

_BinOpStrs = {
    "add": "+",
    "mul": "*",
    "sub": "-",
    "floordiv": "//",
    "truediv": "/",
    "truediv_left": " /l ",
    "ceildiv": "/",
    "ceildiv_left": " /l ",
}


def _math_get_dim_via_bin_op(dims: Sequence[Union[Dim, int]], op_kind: str) -> Dim:
    dims = [d if isinstance(d, _d.Dim) else _make_constant_static_dim(d) for d in dims]
    if all([d.dimension is not None for d in dims]):
        op = _BinOps[op_kind]
        dim_value = dims[0].dimension
        for d in dims[1:]:
            dim_value = op(dim_value, d.dimension)
    else:
        dim_value = None
    if all([d.is_constant_static_dim() for d in dims]):
        return _make_constant_static_dim(dim_value, kind=_get_merged_dim_kind(dims))
    desc = _BinOpStrs[op_kind].join([_get_description(d) for d in dims])
    if op_kind.startswith("ceildiv"):
        desc = f"⌈{desc}⌉"
    return _d.Dim(
        kind=_get_merged_dim_kind(dims),
        description=desc,
        dimension=dim_value,
        derived_from_op=Op(kind=op_kind, inputs=list(dims)),
        derived_from_tag=_representative_tag(dims),
    )


def _is_const_dim_value(d: Union[Dim, int], value: int) -> bool:
    if isinstance(d, int):
        return d == value
    elif isinstance(d, _d.Dim):
        return d.is_constant_static_dim() and d.dimension == value
    else:
        raise TypeError(f"unexpected type {type(d)}")


class _MathFindMatchingAdditive:
    def __init__(self, start: Dim, *, max_depth: int = 2, right: bool, other: Union[int, Dim]):
        self.start = start
        self.max_depth = max_depth
        self.right = right
        self.other = other

    def _check_and_maybe_replace(self, candidate: Dim) -> Optional[Dim]:
        """
        Check and return potential replacement for candidate, when adding `other` to it.
        """
        other = self.other
        if isinstance(other, int) or other.is_constant_static_dim():
            if candidate.is_constant_static_dim():
                return _math_get_dim_via_bin_op([candidate, other] if self.right else [other, candidate], "add")
            return None
        if candidate == other:
            return candidate.__rmul__(2)
        c_op = candidate.derived_from_op
        if c_op and c_op.kind == "mul" and len(c_op.inputs) == 2 and c_op.inputs[1] == other:
            factor = (c_op.inputs[0] + 1) if self.right else (1 + c_op.inputs[0])
            if factor.is_constant_static_dim():
                if factor.dimension == 0:
                    return factor
                factor = factor.dimension
            return factor * c_op.inputs[1]
        o_op = other.derived_from_op
        if not o_op or o_op.kind != "mul" or len(o_op.inputs) != 2:
            return None
        o_base, other = o_op.inputs  # continue checking this
        if candidate == other:
            factor = (1 + o_base) if self.right else (o_base + 1)
            if factor.is_constant_static_dim():
                if factor.dimension == 0:
                    return factor
                factor = factor.dimension
            return factor * candidate
        if c_op and c_op.kind == "mul" and len(c_op.inputs) == 2 and c_op.inputs[1] == other:
            factor = (c_op.inputs[0] + o_base) if self.right else (o_base + c_op.inputs[0])
            if factor.is_constant_static_dim():
                if factor.dimension == 0:
                    return factor
                factor = factor.dimension
            return factor * c_op.inputs[1]
        return None

    def search_and_maybe_replace(self) -> Optional[Dim]:
        """search"""
        cur = self.start
        depth = 0
        history = []
        while True:
            res_cur = self._check_and_maybe_replace(cur)
            if res_cur:
                if depth > 0 and res_cur.is_constant_static_dim() and res_cur.dimension == 0:
                    res_cur = history.pop(-1)
                res = res_cur
                for h in reversed(history):
                    res = _math_get_dim_via_bin_op([h, res] if self.right else [res, h], "add")
                return res
            depth += 1
            if depth > self.max_depth:
                return None
            op = cur.derived_from_op
            if not op or op.kind != "add" or len(op.inputs) != 2:
                return None
            cur = op.inputs[1 if self.right else 0]
            hist = op.inputs[0 if self.right else 1]
            history.append(hist)


def _math_find_matching_mult(start: Dim, other: Union[int, Dim], *, right: bool) -> Optional[Dim]:
    # we assume, if other is Dim, then it is not constant static dim
    if isinstance(other, int) and start.is_constant_static_dim():
        return _math_get_dim_via_bin_op([start, other] if right else [other, start], "mul")
    c_op = start.derived_from_op
    if c_op and c_op.kind == "mul" and len(c_op.inputs) == 2:
        if right:
            return c_op.inputs[0] * (c_op.inputs[1] * other)
        # Don't do right=False -> (other * c_op.inputs[0]) * c_op.inputs[1],
        # because this can lead to infinite recursions,
        # and also we don't have a proper normalized form for multiplication.
        # However, if both left-most factors are constant static dims, then we can merge it.
        elif isinstance(other, int) and c_op.inputs[0].is_constant_static_dim():
            return (other * c_op.inputs[0].dimension) * c_op.inputs[1]
    return None


_DivKindToMeth: Dict[str, Callable[[Dim, Dim], Dim]] = {
    "truediv": _DimMixin.div_right,
    "truediv_left": _DimMixin.div_left,
    "ceildiv": _DimMixin.ceildiv_right,
    "ceildiv_left": _DimMixin.ceildiv_left,
    "floordiv": _DimMixin.__floordiv__,
}


def _math_find_matching_div(start: Dim, other: Union[int, Dim], *, right: bool, kind: str) -> Optional[Dim]:
    if (isinstance(other, int) or other.is_constant_static_dim()) and start.is_constant_static_dim():
        return _math_get_dim_via_bin_op([start, other] if right else [other, start], kind)
    c_op = start.derived_from_op
    if c_op and c_op.kind == kind and len(c_op.inputs) == 2:
        meth = _DivKindToMeth[kind]
        return meth(c_op.inputs[0], c_op.inputs[1] * other if right else other * c_op.inputs[1])
    return None


class Op:
    """
    Op on :class:`Dim` which results in a derived :class:`Dim`.
    """

    def __init__(self, kind: str, inputs: List[Dim], attribs: Optional[Dict[str, Any]] = None):
        """
        :param kind: "add", "sub", "mul", "ceildiv"
        :param inputs:
        :param 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):
        with util.guard_infinite_recursion(Op.__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 "?")


def _get_merged_dim_kind(dim_tags: Sequence[Dim]) -> 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 _representative_tag(terms: Sequence[Dim]) -> Optional[Dim]:
    if any([not term_.auto_generated for term_ in terms]):
        # Always prefer non-auto-generated.
        terms = [term_ for term_ in terms if not term_.auto_generated]
    # First find any dynamic.
    for term_ in terms:
        if term_.is_dynamic_seq_length():
            return term_
    # Now find non-unspecified.
    for term_ in terms:
        if term_.kind != DimTypes.Unspecified:
            return term_
    # Now find any.
    for term_ in terms:
        return term_
    return None


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


class _CacheDimMath:
    """op (add,sub,...), operand -> Dim"""

    class _OperandCache:
        def __init__(self):
            self.dims: MutableMapping[Dim, Dim] = weakref.WeakKeyDictionary()
            self.statics: Dict[int, Dim] = {}

    def __init__(self):
        self._ops: Dict[str, _CacheDimMath._OperandCache] = {}

    def __repr__(self):
        return "_CacheDimMath({%s})" % ", ".join("%r: %r" % (k, v) for k, v in self.items())

    def _get_op_dict(self, __key: Tuple[str, Union[Dim, int]]) -> _OperandCache:
        if __key[0] in self._ops:
            return self._ops[__key[0]]
        else:
            op_dict = self._OperandCache()
            self._ops[__key[0]] = op_dict
            return op_dict

    def __setitem__(self, __key: Tuple[str, Union[Dim, int]], __value: Dim):
        op_dict = self._get_op_dict(__key)
        if isinstance(__key[1], int):
            value_dict = op_dict.statics
        else:
            value_dict = op_dict.dims
        if __key[1] in value_dict:
            value_dict[__key[1]] = __value
            return
        if len(value_dict) >= 5:
            # Just to avoid memory leaks.
            value_dict.clear()
        value_dict[__key[1]] = __value

    def __delitem__(self, __key: Tuple[str, Union[Dim, int]]):
        op_dict = self._ops[__key[0]]
        if isinstance(__key[1], int):
            del op_dict.statics[__key[1]]
        else:
            del op_dict.dims[__key[1]]

    def __getitem__(self, __key: Tuple[str, Union[Dim, int]]) -> Dim:
        op_dict = self._ops[__key[0]]
        if isinstance(__key[1], int):
            return op_dict.statics[__key[1]]
        else:
            return op_dict.dims[__key[1]]

    def __contains__(self, __key: Tuple[str, Union[Dim, int]]) -> bool:
        op_dict = self._ops.get(__key[0])
        if not op_dict:
            return False
        if isinstance(__key[1], int):
            return __key[1] in op_dict.statics
        else:
            return __key[1] in op_dict.dims

    def get(self, __key: Tuple[str, Union[Dim, int]], default: Optional[Dim] = None) -> Optional[Dim]:
        """get"""
        op_dict = self._ops.get(__key[0])
        if not op_dict:
            return default
        if isinstance(__key[1], int):
            return op_dict.statics.get(__key[1], default)
        else:
            return op_dict.dims.get(__key[1], default)

    def setdefault(self, __key: Tuple[str, Union[Dim, int]], __value: Dim):
        """setdefault"""
        existing = self.get(__key)
        if existing is not None:
            return existing
        self[__key] = __value
        return __value

    def clear(self):
        """clear"""
        self._ops.clear()

    def clear_dynamic(self):
        """clear dynamic part"""
        for op_dict in self._ops.values():
            for k, v in list(op_dict.dims.items()):
                if v.dyn_size_ext is not None or v.dimension is None:
                    del op_dict.dims[k]

    def __len__(self):
        count = 0
        for op_dict in self._ops.values():
            count += len(op_dict.statics)
            count += len(op_dict.dims)
        return count

    def items(self):
        """items"""
        for op_name, op_dict in self._ops.items():
            for key, value in op_dict.statics.items():
                yield (op_name, key), value
            for key, value in op_dict.dims.items():
                yield (op_name, key), value

    def keys(self):
        """keys"""
        for k, v in self.items():
            yield k

    def values(self):
        """values"""
        for k, v in self.items():
            yield v

    def __iter__(self):
        yield from self.keys()


def _behavior_version_reset_callback():
    # Reset things we did in _handle_new_min_version.
    _DimMixin._SimpleEquality = False
    # noinspection PyProtectedMember
    _DimMixin.__eq__ = _DimMixin.is_equal
    # noinspection PyProtectedMember
    _DimMixin.__ne__ = _DimMixin._ne_generic


def _behavior_version_handle_new_min_version_callback():
    if util.BehaviorVersion.get() >= 16:
        _DimMixin._SimpleEquality = True
        # noinspection PyProtectedMember
        _DimMixin.__eq__ = _DimMixin._eq_simple
        # noinspection PyProtectedMember
        _DimMixin.__ne__ = _DimMixin._ne_simple


def _setup():
    util.BehaviorVersion.reset_callbacks.append(_behavior_version_reset_callback)
    util.BehaviorVersion.handle_new_min_version_callbacks.append(_behavior_version_handle_new_min_version_callback)
    _behavior_version_handle_new_min_version_callback()


_setup()
