import re
import warnings
from numbers import Number
from typing import Tuple, Callable, List, Union, Any, Sequence, Optional, Dict

from .. import math


BATCH_DIM = 'batch'
SPATIAL_DIM = 'spatial'
CHANNEL_DIM = 'channel'
INSTANCE_DIM = 'înstance'
DUAL_DIM = 'dual'

SUPERSCRIPT = {SPATIAL_DIM: "ˢ", CHANNEL_DIM: "ᶜ", INSTANCE_DIM: "ⁱ", BATCH_DIM: "ᵇ", DUAL_DIM: "ᵈ", None: "⁻"}  # ᵃᵇᶜᵈᵉᶠᵍʰⁱʲᵏˡᵐⁿᵒᵖʳˢᵗᵘᵛʷˣʸᶻ
CHAR = {SPATIAL_DIM: "s", CHANNEL_DIM: "c", INSTANCE_DIM: "i", BATCH_DIM: "b", DUAL_DIM: "d", None: "-"}  # ᵃᵇᶜᵈᵉᶠᵍʰⁱʲᵏˡᵐⁿᵒᵖʳˢᵗᵘᵛʷˣʸᶻ
INV_CHAR = {v: k for k, v in CHAR.items()}

DEBUG_CHECKS = []

DimFilter = Union[str, Sequence, set, 'Shape', Callable, None]
try:
    DimFilter.__doc__ = """Dimension filters can be used with `Shape.only()` and `Shype.without()`, making them the standard tool for specifying sets of dimensions.

    The following types can be used as dimension filters:

    * `Shape` instances
    * `tuple` or `list` objects containing dimension names as `str`
    * Single `str` listing comma-separated dimension names
    * Any function `filter(Shape) -> Shape`, such as `math.batch()`, `math.non_batch()`, `math.spatial()`, etc.
    """  # docstring must be set explicitly
except AttributeError:  # on older Python versions, this is not possible
    pass


def enable_debug_checks():
    """
    Once called, additional type checks are enabled.
    This may result in a noticeable drop in performance.
    """
    DEBUG_CHECKS.append(True)


class Shape:
    """
    Shapes enumerate dimensions, each consisting of a name, size and type.

    There are five types of dimensions: `batch`, `dual`, `spatial`, `channel`, and `instance`.
    """

    def __init__(self, sizes: tuple, names: tuple, types: tuple, item_names: tuple):
        """
        To construct a `Shape`, use `batch`, `dual`, `spatial`, `channel` or `instance`, depending on the desired dimension type.
        To create a shape with multiple types, use `merge_shapes()`, `concat_shapes()` or the syntax `shape1 & shape2`.

        The `__init__` constructor is for internal use only.
        """
        if len(sizes) > 0 and any(s is not None and not isinstance(s, int) for s in sizes):
            from ._tensors import Tensor
            sizes = tuple([s if isinstance(s, Tensor) or s is None else int(s) for s in sizes])  # TODO replace this by an assert
        self.sizes: tuple = sizes
        """
        Ordered dimension sizes as `tuple`.
        The size of a dimension can be an `int` or a `Tensor` for [non-uniform shapes](https://tum-pbs.github.io/PhiML/Non_Uniform.html).
        
        See Also:
            `Shape.get_size()`, `Shape.size`, `Shape.shape`.
        """
        self.names: Tuple[str, ...] = names
        """
        Ordered dimension names as `tuple[str]`.
        
        See Also:
            `Shape.name`.
        """
        self.types: Tuple[str, ...] = types  # undocumented, may be private
        self.item_names: Tuple[Optional[Tuple[str, ...]], ...] = (None,) * len(sizes) if item_names is None else item_names  # undocumented
        if DEBUG_CHECKS:
            assert len(sizes) == len(names) == len(types) == len(item_names), f"sizes={sizes}, names={names}, types={types}, item_names={item_names}"
            assert len(set(names)) == len(names), f"Duplicate dimension names: {names}"
            assert all(isinstance(n, str) for n in names), f"All names must be of type string but got {names}"
            assert isinstance(self.item_names, tuple)
            assert all([items is None or isinstance(items, tuple) for items in self.item_names])
            assert all([items is None or all([isinstance(n, str) for n in items]) for items in self.item_names])
            from ._tensors import Tensor
            for name, size in zip(names, sizes):
                if isinstance(size, Tensor):
                    assert size.rank > 0
            for name, size, item_names in zip(self.names, self.sizes, self.item_names):
                if item_names is not None:
                    try:
                        int(size)
                    except Exception:
                        raise AssertionError(f"When item names are present, the size must be an integer type")
                    assert len(item_names) == size, f"Number of item names ({len(item_names)}) does not match size {size}"
                    for item_name in item_names:
                        assert item_name, f"Empty item name"
                    assert len(set(item_names)) == len(item_names), f"Duplicate item names in shape {self} at dim '{name}': {item_names}"
            for name, type in zip(names, types):
                if type == DUAL_DIM:
                    assert name.startswith('~'), f"Dual dimensions must start with '~' but got '{name}' in {self}"

    def _check_is_valid_tensor_shape(self):
        if DEBUG_CHECKS:
            from ._tensors import Tensor
            for name, size in zip(self.names, self.sizes):
                if size is not None and isinstance(size, Tensor):
                    assert size.rank > 0
                    for dim in size.shape.names:
                        assert dim in self.names, f"Dimension {name} varies along {dim} but {dim} is not part of the Shape {self}"

    def _to_dict(self, include_sizes=True):
        result = dict(names=self.names, types=self.types, item_names=self.item_names)
        if include_sizes:
            if not all([isinstance(s, int)] for s in self.sizes):
                raise NotImplementedError()
            result['sizes'] = self.sizes
        return result

    @staticmethod
    def _from_dict(dict_: dict):
        names = tuple(dict_['names'])
        sizes = list(dict_['sizes']) if 'sizes' in dict_ else [None] * len(names)
        item_names = tuple([None if n is None else tuple(n) for n in dict_['item_names']])
        for i, n in enumerate(item_names):
            if n and sizes[i] is None:
                sizes[i] = len(n)
        return Shape(tuple(sizes), names, tuple(dict_['types']), item_names)

    @property
    def name_list(self):
        return list(self.names)

    @property
    def _named_sizes(self):
        return zip(self.names, self.sizes)

    @property
    def _dimensions(self):
        return zip(self.sizes, self.names, self.types, self.item_names)

    @property
    def untyped_dict(self):
        """
        Returns:
            `dict` containing dimension names as keys.
                The values are either the item names as `tuple` if available, otherwise the size.
        """
        return {name: self.get_item_names(i) or self.get_size(i) for i, name in enumerate(self.names)}

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

    def __contains__(self, item):
        if isinstance(item, (str, tuple, list)):
            dims = parse_dim_order(item)
            return all(dim in self.names for dim in dims)
        elif isinstance(item, Shape):
            return all([d in self.names for d in item.names])
        else:
            raise ValueError(item)

    def isdisjoint(self, other: Union['Shape', tuple, list, str]):
        """ Shapes are disjoint if all dimension names of one shape do not occur in the other shape. """
        other = parse_dim_order(other)
        return not any(dim in self.names for dim in other)

    def __iter__(self):
        return iter(self[i] for i in range(self.rank))

    def index(self, dim: Union[str, 'Shape', None]) -> int:
        """
        Finds the index of the dimension within this `Shape`.

        See Also:
            `Shape.indices()`.

        Args:
            dim: Dimension name or single-dimension `Shape`.

        Returns:
            Index as `int`.
        """
        if dim is None:
            return None
        elif isinstance(dim, str):
            if dim not in self.names:
                raise ValueError(f"Shape {self} has no dimension '{dim}'")
            return self.names.index(dim)
        elif isinstance(dim, Shape):
            assert dim.rank == 1, f"index() requires a single dimension as input but got {dim}. Use indices() for multiple dimensions."
            return self.names.index(dim.name)
        else:
            raise ValueError(f"index() requires a single dimension as input but got {dim}")

    def indices(self, dims: Union[tuple, list, 'Shape']) -> Tuple[int]:
        """
        Finds the indices of the given dimensions within this `Shape`.

        See Also:
            `Shape.index()`.

        Args:
            dims: Sequence of dimensions as `tuple`, `list` or `Shape`.

        Returns:
            Indices as `tuple[int]`.
        """
        if isinstance(dims, (list, tuple, set)):
            return tuple([self.index(n) for n in dims if n in self.names])
        elif isinstance(dims, Shape):
            return tuple([self.index(n) for n in dims.names if n in self.names])
        else:
            raise ValueError(f"indices() requires a sequence of dimensions but got {dims}")

    def get_size(self, dim: Union[str, 'Shape', int], default=None):
        """
        See Also:
            `Shape.get_sizes()`, `Shape.size`

        Args:
            dim: Dimension, either as name `str` or single-dimension `Shape` or index `int`.
            default: (Optional) If the dim does not exist, return this value instead of raising an error.

        Returns:
            Size associated with `dim` as `int` or `Tensor`.
        """
        if isinstance(dim, int):
            assert default is None, "Cannot use a default value when passing an int for dim"
            return self.sizes[dim]
        if isinstance(dim, Shape):
            assert dim.rank == 1, f"get_size() requires a single dimension but got {dim}. Use indices() to get multiple sizes."
            dim = dim.name
        if isinstance(dim, str):
            if dim not in self.names:
                if default is None:
                    raise KeyError(f"get_size() failed because '{dim}' is not part of Shape {self} and no default value was provided")
                else:
                    return default
            return self.sizes[self.names.index(dim)]
        else:
            raise ValueError(f"get_size() requires a single dim name but got {dim}. Use indices() to get multiple sizes.")

    def get_sizes(self, dims: Union[tuple, list, 'Shape']) -> tuple:
        """
        See Also:
            `Shape.get_size()`

        Args:
            dims: Dimensions as `tuple`, `list` or `Shape`.

        Returns:
            `tuple`
        """
        assert isinstance(dims, (tuple, list, Shape)), f"get_sizes() requires a sequence of dimensions but got {dims}"
        return tuple([self.get_size(dim) for dim in dims])

    def get_type(self, dim: Union[str, 'Shape']) -> str:
        # undocumented, use get_dim_type() instead.
        if isinstance(dim, str):
            return self.types[self.names.index(dim)]
        elif isinstance(dim, Shape):
            assert dim.rank == 1, f"Shape.get_type() only accepts single-dimension Shapes but got {dim}"
            return self.types[self.names.index(dim.name)]
        else:
            raise ValueError(dim)

    def get_dim_type(self, dim: Union[str, 'Shape']) -> Callable:
        """
        Args:
            dim: Dimension, either as name `str` or single-dimension `Shape`.

        Returns:
            Dimension type, one of `batch`, `spatial`, `instance`, `channel`.
        """
        return DIM_FUNCTIONS[self.get_type(dim)]

    def get_types(self, dims: Union[tuple, list, 'Shape']) -> tuple:
        # undocumented, do not use
        if isinstance(dims, (tuple, list)):
            return tuple(self.get_type(n) for n in dims)
        elif isinstance(dims, Shape):
            return tuple(self.get_type(n) for n in dims.names)
        else:
            raise ValueError(dims)

    def get_item_names(self, dim: Union[str, 'Shape', int], fallback_spatial=False) -> Union[tuple, None]:
        """
        Args:
            fallback_spatial: If `True` and no item names are defined for `dim` and `dim` is a channel dimension, the spatial dimension names are interpreted as item names along `dim` in the order they are listed in this `Shape`.
            dim: Dimension, either as `int` index, `str` name or single-dimension `Shape`.

        Returns:
            Item names as `tuple` or `None` if not defined.
        """
        if isinstance(dim, int):
            result = self.item_names[dim]
        elif isinstance(dim, str):
            result = self.item_names[self.index(dim)]
        elif isinstance(dim, Shape):
            assert dim.rank == 1, f"Shape.get_type() only accepts single-dimension Shapes but got {dim}"
            result = self.item_names[self.names.index(dim.name)]
        else:
            raise ValueError(dim)
        if result is not None:
            return result
        elif fallback_spatial and self.spatial_rank == self.get_size(dim) and self.get_type(dim) == CHANNEL_DIM:
            return self.spatial.names
        else:
            return None

    def flipped(self, dims: Union[List[str], Tuple[str]]):
        item_names = list(self.item_names)
        for dim in dims:
            if dim in self.names:
                dim_i_n = self.get_item_names(dim)
                if dim_i_n is not None:
                    item_names[self.index(dim)] = tuple(reversed(dim_i_n))
        return Shape(self.sizes, self.names, self.types, tuple(item_names))

    def __getitem__(self, selection):
        if isinstance(selection, int):
            return Shape((self.sizes[selection],), (self.names[selection],), (self.types[selection],), (self.item_names[selection],))
        elif isinstance(selection, slice):
            return Shape(self.sizes[selection], self.names[selection], self.types[selection], self.item_names[selection])
        elif isinstance(selection, str):
            if ',' in selection:
                selection = [self.index(s.strip()) for s in selection.split(',')]
            else:
                selection = self.index(selection)
            return self[selection]
        elif isinstance(selection, Shape):
            selection = selection.names
        if isinstance(selection, (tuple, list)):
            selection = [self.index(s) if isinstance(s, str) else s for s in selection]
            return Shape(tuple([self.sizes[i] for i in selection]), tuple([self.names[i] for i in selection]), tuple([self.types[i] for i in selection]), tuple([self.item_names[i] for i in selection]))
        raise AssertionError("Can only access shape elements as shape[int], shape[str], shape[slice], shape[Sequence] or shape[Shape]")

    @property
    def reversed(self):
        return Shape(tuple(reversed(self.sizes)), tuple(reversed(self.names)), tuple(reversed(self.types)), tuple(reversed(self.item_names)))

    @property
    def batch(self) -> 'Shape':
        """
        Filters this shape, returning only the batch dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t == BATCH_DIM]]

    @property
    def non_batch(self) -> 'Shape':
        """
        Filters this shape, returning only the non-batch dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t != BATCH_DIM]]

    @property
    def spatial(self) -> 'Shape':
        """
        Filters this shape, returning only the spatial dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t == SPATIAL_DIM]]

    @property
    def non_spatial(self) -> 'Shape':
        """
        Filters this shape, returning only the non-spatial dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t != SPATIAL_DIM]]

    @property
    def instance(self) -> 'Shape':
        """
        Filters this shape, returning only the instance dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t == INSTANCE_DIM]]

    @property
    def non_instance(self) -> 'Shape':
        """
        Filters this shape, returning only the non-instance dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t != INSTANCE_DIM]]

    @property
    def channel(self) -> 'Shape':
        """
        Filters this shape, returning only the channel dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t == CHANNEL_DIM]]

    @property
    def non_channel(self) -> 'Shape':
        """
        Filters this shape, returning only the non-channel dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t != CHANNEL_DIM]]

    @property
    def dual(self) -> 'Shape':
        """
        Filters this shape, returning only the dual dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t == DUAL_DIM]]

    @property
    def non_dual(self) -> 'Shape':
        """
        Filters this shape, returning only the non-dual dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t != DUAL_DIM]]

    @property
    def primal(self) -> 'Shape':
        """
        Filters this shape, returning only the dual dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t not in [DUAL_DIM, BATCH_DIM]]]

    @property
    def non_primal(self) -> 'Shape':
        """
        Filters this shape, returning only batch and dual dimensions as a new `Shape` object.

        See also:
            `Shape.batch`, `Shape.spatial`, `Shape.instance`, `Shape.channel`, `Shape.dual`, `Shape.non_batch`, `Shape.non_spatial`, `Shape.non_instance`, `Shape.non_channel`, `Shape.non_dual`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, t in enumerate(self.types) if t in [DUAL_DIM, BATCH_DIM]]]

    @property
    def transposed(self):
        if self.channel_rank > 0:
            replacement = {DUAL_DIM: CHANNEL_DIM, CHANNEL_DIM: DUAL_DIM}
        elif self.instance_rank > 0:
            replacement = {DUAL_DIM: INSTANCE_DIM, INSTANCE_DIM: DUAL_DIM}
        elif self.spatial_rank > 0:
            replacement = {DUAL_DIM: SPATIAL_DIM, SPATIAL_DIM: DUAL_DIM}
        elif self.dual_rank > 0:
            warnings.warn(f"Transposing {self} is ill-defined because there are not primal dims. Replacing dual dims by channel dims.", SyntaxWarning)
            replacement = {DUAL_DIM: CHANNEL_DIM}
        else:
            raise ValueError(f"Cannot transpose shape {self} as it has no channel or instance or spatial dims.")
        return self._with_types(tuple([replacement.get(t, t) for t in self.types]))

    def transpose(self, dims: DimFilter):
        if callable(dims) and dims in TYPE_BY_FUNCTION:
            dims = TYPE_BY_FUNCTION[dims]
            replacement = {DUAL_DIM: dims, dims: DUAL_DIM}
            return self._with_types(tuple([replacement.get(t, t) for t in self.types]))
        dims = self.only(dims)
        return self.replace(dims, dims.transposed)

    @property
    def non_singleton(self) -> 'Shape':
        """
        Filters this shape, returning only non-singleton dimensions as a new `Shape` object.
        Dimensions are singleton if their size is exactly `1`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, s in enumerate(self.sizes) if not _size_equal(s, 1)]]

    @property
    def singleton(self) -> 'Shape':
        """
        Filters this shape, returning only singleton dimensions as a new `Shape` object.
        Dimensions are singleton if their size is exactly `1`.

        Returns:
            New `Shape` object
        """
        return self[[i for i, s in enumerate(self.sizes) if _size_equal(s, 1)]]

    def assert_all_sizes_defined(self):
        """
        Filters this shape, returning only singleton dimensions as a new `Shape` object.
        Dimensions are singleton if their size is exactly `1`.

        Returns:
            New `Shape` object
        """
        for n, s in zip(self.names, self.sizes):
            assert s is not None, f"All sizes must be defined but dim '{n}' is undefined in shape {self}"

    def as_channel(self):
        """Returns a copy of this `Shape` with all dimensions of type *channel*."""
        return channel(**self.untyped_dict)

    def as_batch(self):
        """Returns a copy of this `Shape` with all dimensions of type *batch*."""
        return batch(**self.untyped_dict)

    def as_spatial(self):
        """Returns a copy of this `Shape` with all dimensions of type *spatial*."""
        return spatial(**self.untyped_dict)

    def as_instance(self):
        """Returns a copy of this `Shape` with all dimensions of type *instance*."""
        return instance(**self.untyped_dict)

    def as_dual(self):
        """Returns a copy of this `Shape` with all dimensions of type *dual*."""
        return dual(**self.untyped_dict)

    def as_type(self, new_type: Callable):
        """Returns a copy of this `Shape` with all dimensions of the given type, either `batch`, `dual`, `spatial`, `instance`, or `channel` ."""
        return new_type(**self.untyped_dict)

    def _more_dual(self):
        return Shape(self.sizes, tuple('~' + n for n in self.names), (DUAL_DIM,) * len(self.names), self.item_names)

    def _less_dual(self, default_type='unknown_primal'):
        names = tuple(n[1:] if n.startswith('~') else n for n in self.names)
        types = [t if t != DUAL_DIM else (DUAL_DIM if n.startswith('~~') else default_type) for n, t in zip(self.names, self.types)]
        return Shape(self.sizes, names, tuple(types), self.item_names)

    def unstack(self, dim='dims') -> Tuple['Shape']:
        """
        Slices this `Shape` along a dimension.
        The dimension listing the sizes of the shape is referred to as `'dims'`.

        Non-uniform tensor shapes may be unstacked along other dimensions as well, see
        https://tum-pbs.github.io/PhiML/Non_Uniform.html

        Args:
            dim: dimension to unstack

        Returns:
            slices of this shape
        """
        if dim == 'dims':
            return tuple(Shape((self.sizes[i],), (self.names[i],), (self.types[i],), (self.item_names[i],)) for i in range(self.rank))
        if dim not in self and self.is_uniform:
            return tuple([self])
        from ._tensors import Tensor
        if dim in self:
            inner = self.without(dim)
            dim_size = self.get_size(dim)
        else:
            inner = self
            dim_size = self.shape.get_size(dim)
        sizes = []
        for size in inner.sizes:
            if isinstance(size, Tensor) and dim in size.shape:
                sizes.append(size._unstack(dim))
                dim_size = size.shape.get_size(dim)
            else:
                sizes.append(size)
        assert isinstance(dim_size, int)
        shapes = tuple(Shape(tuple([int(size[i]) if isinstance(size, tuple) else size for size in sizes]), inner.names, inner.types, inner.item_names) for i in range(dim_size))
        return shapes

    @property
    def name(self) -> str:
        """
        Only for Shapes containing exactly one single dimension.
        Returns the name of the dimension.

        See Also:
            `Shape.names`.
        """
        assert self.rank == 1, f"Shape.name is only defined for shapes of rank 1. shape={self}"
        return self.names[0]

    @property
    def size(self):
        """
        Only for Shapes containing exactly one single dimension.
        Returns the size of the dimension.

        See Also:
            `Shape.sizes`, `Shape.get_size()`.
        """
        assert self.rank == 1, f"Shape.size is only defined for shapes of rank 1 but has dims {self}"
        return self.sizes[0]

    @property
    def type(self) -> str:
        """
        Only for Shapes containing exactly one single dimension.
        Returns the type of the dimension.

        See Also:
            `Shape.get_type()`.
        """
        assert self.rank == 1, "Shape.type is only defined for shapes of rank 1."
        return self.types[0]

    @property
    def dim_type(self):
        types = set(self.types)
        assert len(types) == 1, f"Shape contains multiple types: {self}"
        return DIM_FUNCTIONS[next(iter(types))]

    def __int__(self):
        assert self.rank == 1, "int(Shape) is only defined for shapes of rank 1."
        return self.sizes[0]

    def mask(self, names: Union[tuple, list, set, 'Shape']):
        """
        Returns a binary sequence corresponding to the names of this Shape.
        A value of 1 means that a dimension of this Shape is contained in `names`.

        Args:
          names: instance of dimension
          names: tuple or list or set: 

        Returns:
          binary sequence

        """
        if isinstance(names, str):
            names = [names]
        elif isinstance(names, Shape):
            names = names.names
        mask = [1 if name in names else 0 for name in self.names]
        return tuple(mask)

    def __repr__(self):
        def size_repr(size, items):
            if items is not None:
                items_str = ",".join(items)
                return items_str if len(items_str) <= 12 else f"{size}:{items[0][:5]}..."
            return size

        strings = [f"{name}{SUPERSCRIPT.get(dim_type, '?')}={size_repr(size, items)}" for size, name, dim_type, items in self._dimensions]
        return '(' + ', '.join(strings) + ')'

    def __eq__(self, other):
        if not isinstance(other, Shape):
            return False
        if self.names != other.names or self.types != other.types:
            return False
        for size1, size2 in zip(self.sizes, other.sizes):
            equal = size1 == size2
            assert isinstance(equal, (bool, math.Tensor))
            if isinstance(equal, math.Tensor):
                equal = equal.all
            if not equal:
                return False
        for names1, names2 in zip(self.item_names, other.item_names):
            if names1 != names2:
                return False
        return True

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

    def __bool__(self):
        return self.rank > 0

    def _reorder(self, names: Union[tuple, list, 'Shape']) -> 'Shape':
        assert len(names) == self.rank
        if isinstance(names, Shape):
            names = names.names
        order = [self.index(n) for n in names]
        return self[order]

    def _order_group(self, names: Union[tuple, list, 'Shape']) -> list:
        """ Reorders the dimensions of this `Shape` so that `names` are clustered together and occur in the specified order. """
        if isinstance(names, Shape):
            names = names.names
        result = []
        for dim in self.names:
            if dim not in result:
                if dim in names:
                    result.extend(names)
                else:
                    result.append(dim)
        return result

    def __and__(self, other):
        if other is dual:
            return concat_shapes(self, self.primal.as_dual())
        return merge_shapes(self, other)

    def __rand__(self, other):
        if other is dual:
            return concat_shapes(self.primal.as_dual(), self)
        return merge_shapes(self, other)

    def _expand(self, dim: 'Shape', pos=None) -> 'Shape':
        """**Deprecated.** Use `phiml.math.merge_shapes()` or `phiml.math.concat_shapes()` instead. """
        if not dim:
            return self
        assert dim.name not in self, f"Cannot expand shape {self} by {dim} because dimension already exists."
        assert isinstance(dim, Shape) and dim.rank == 1, f"Shape.expand() requires a single dimension as a Shape but got {dim}"
        if pos is None:
            same_type_dims = self[[i for i, t in enumerate(self.types) if t == dim.type]]
            if len(same_type_dims) > 0:
                pos = self.index(same_type_dims.names[0])
            else:
                pos = {
                    BATCH_DIM: 0,
                    DUAL_DIM: self.batch_rank,
                    INSTANCE_DIM: self.batch_rank + self.dual_rank,
                    SPATIAL_DIM: self.batch.rank + self.dual_rank + self.instance_rank,
                    CHANNEL_DIM: self.rank + 1
                }[dim.type]
        elif pos < 0:
            pos += self.rank + 1
        sizes = list(self.sizes)
        names = list(self.names)
        types = list(self.types)
        item_names = list(self.item_names)
        sizes.insert(pos, dim.size)
        names.insert(pos, dim.name)
        types.insert(pos, dim.type)
        item_names.insert(pos, dim.item_names[0])
        return Shape(tuple(sizes), tuple(names), tuple(types), tuple(item_names))

    def without(self, dims: 'DimFilter') -> 'Shape':
        """
        Builds a new shape from this one that is missing all given dimensions.
        Dimensions in `dims` that are not part of this Shape are ignored.
        
        The complementary operation is `Shape.only()`.

        Args:
          dims: Single dimension (str) or instance of dimensions (tuple, list, Shape)
          dims: Dimensions to exclude as `str` or `tuple` or `list` or `Shape`. Dimensions that are not included in this shape are ignored.

        Returns:
          Shape without specified dimensions
        """
        if dims is None:  # subtract none
            return self
        elif callable(dims):
            dims = dims(self)
        if isinstance(dims, str):
            return self[[i for i in range(self.rank) if self.names[i] not in parse_dim_order(dims)]]
        elif isinstance(dims, Shape):
            return self[[i for i in range(self.rank) if self.names[i] not in dims.names]]
        if isinstance(dims, (tuple, list, set)) and all([isinstance(d, str) for d in dims]):
            return self[[i for i in range(self.rank) if self.names[i] not in dims]]
        elif isinstance(dims, (tuple, list, set)):
            result = self
            for wo in dims:
                result = result.without(wo)
            return result
        else:
            raise ValueError(dims)

    def only(self, dims: 'DimFilter', reorder=False):
        """
        Builds a new shape from this one that only contains the given dimensions.
        Dimensions in `dims` that are not part of this Shape are ignored.
        
        The complementary operation is :func:`Shape.without`.

        Args:
          dims: comma-separated dimension names (str) or instance of dimensions (tuple, list, Shape) or filter function.
          reorder: If `False`, keeps the dimension order as defined in this shape.
            If `True`, reorders the dimensions of this shape to match the order of `dims`.

        Returns:
          Shape containing only specified dimensions

        """
        if dims is None:  # keep none
            return EMPTY_SHAPE
        if callable(dims):
            dims = dims(self)
        if isinstance(dims, str):
            dims = parse_dim_order(dims)
        if isinstance(dims, Shape):
            dims = dims.names
        if isinstance(dims, (tuple, list, set)):
            if all(isinstance(d, int) for d in dims):
                if not reorder:
                    dims = tuple(sorted(dims))
                return self[dims]
            dim_names = []
            for d in dims:
                if callable(d):
                    d = d(self)
                if isinstance(d, str):
                    dim_names.append(d)
                elif isinstance(d, Shape):
                    dim_names.extend(d.names)
                else:
                    raise ValueError(f"Format not understood for Shape.only(): {dims}")
            if reorder:
                dim_names = [d.name if isinstance(d, Shape) else d for d in dim_names]
                assert all(isinstance(d, str) for d in dim_names)
                return self[[self.names.index(d) for d in dim_names if d in self.names]]
            else:
                dim_names = [d.name if isinstance(d, Shape) else d for d in dim_names]
                assert all(isinstance(d, str) for d in dim_names)
                return self[[i for i in range(self.rank) if self.names[i] in dim_names]]
        raise ValueError(dims)

    def is_compatible(self, *others: 'Shape'):
        """
        Checks if this shape and the others can be broadcast.

        Args:
            others: Other shapes.

        Returns:
            `True` only if all shapes are compatible.
        """
        try:
            merge_shapes(self, *others)
            return True
        except IncompatibleShapes:
            return False

    @property
    def rank(self) -> int:
        """
        Returns the number of dimensions.
        Equal to `len(shape)`.

        See `Shape.is_empty`, `Shape.batch_rank`, `Shape.spatial_rank`, `Shape.channel_rank`.
        """
        return len(self.sizes)

    @property
    def batch_rank(self) -> int:
        """ Number of batch dimensions """
        return sum([1 for ty in self.types if ty == BATCH_DIM])

    @property
    def instance_rank(self) -> int:
        return sum([1 for ty in self.types if ty == INSTANCE_DIM])

    @property
    def spatial_rank(self) -> int:
        """ Number of spatial dimensions """
        return sum([1 for ty in self.types if ty == SPATIAL_DIM])

    @property
    def dual_rank(self) -> int:
        """ Number of spatial dimensions """
        return sum([1 for ty in self.types if ty == DUAL_DIM])

    @property
    def channel_rank(self) -> int:
        """ Number of channel dimensions """
        return sum([1 for ty in self.types if ty == CHANNEL_DIM])

    @property
    def well_defined(self):
        """
        Returns `True` if no dimension size is `None`.

        Shapes with undefined sizes may be used in `phiml.math.tensor()`, `phiml.math.wrap()`, `phiml.math.stack()` or `phiml.math.concat()`.

        To create an undefined size, call a constructor function (`batch()`, `spatial()`, `channel()`, `instance()`)
        with positional `str` arguments, e.g. `spatial('x')`.
        """
        for size in self.sizes:
            if size is None:
                return False
        return True

    @property
    def defined(self):
        return self[[i for i, size in enumerate(self.sizes) if size is not None]]

    @property
    def undefined(self):
        return self[[i for i, size in enumerate(self.sizes) if size is None]]

    @property
    def shape(self) -> 'Shape':
        """
        Higher-order `Shape`.
        The returned shape will always contain the channel dimension `dims` with a size equal to the `Shape.rank` of this shape.

        For uniform shapes, `Shape.shape` will only contain the dimension `dims` but the shapes of [non-uniform shapes](https://tum-pbs.github.io/PhiML/Non_Uniform.html)
        may contain additional dimensions.

        See Also:
            `Shape.is_uniform`.

        Returns:
            `Shape`.
        """
        from . import Tensor
        shape = Shape((self.rank,), ('dims',), (CHANNEL_DIM,), (self.names,))
        for size in self.sizes:
            if isinstance(size, Tensor):
                shape = shape & size.shape
        return shape

    @property
    def is_uniform(self) -> bool:
        """
        A shape is uniform if it all sizes have a single integer value.

        See Also:
            `Shape.is_non_uniform`, `Shape.shape`.
        """
        from ._tensors import Tensor
        return all(not isinstance(s, Tensor) for s in self.sizes)

    @property
    def is_non_uniform(self) -> bool:
        """
        A shape is non-uniform if the size of any dimension varies along another dimension.

        See Also:
            `Shape.is_uniform`, `Shape.shape`.
        """
        return not self.is_uniform

    @property
    def non_uniform(self) -> 'Shape':
        """
        Returns only the non-uniform dimensions of this shape, i.e. the dimensions whose size varies along another dimension.

        See Also
            `Shape.non_uniform_shape`
        """
        from . import Tensor
        indices = [i for i, size in enumerate(self.sizes) if isinstance(size, Tensor) and size.rank > 0]
        return self[indices]

    @property
    def non_uniform_shape(self):
        """
        Returns the stack dimensions of non-uniform shapes.
        This is equal to `Shape.shape` excluding the `dims` dimension.

        For example, when stacking `(x=3)` and `(x=2)` along `vector`, the resulting shape is non_uniform.
        Its `non_uniform_shape` is `vector` and its `non_uniform` dimension is `x`.

        See Also
            `Shape.non_uniform`.
        """
        from . import Tensor
        shape = EMPTY_SHAPE
        for size in self.sizes:
            if isinstance(size, Tensor):
                shape = shape & size.shape
        return shape

    def with_size(self, size: Union[int, Sequence[str]]):
        """
        Only for single-dimension shapes.
        Returns a `Shape` representing this dimension but with a different size.

        See Also:
            `Shape.with_sizes()`.

        Args:
            size: Replacement size for this dimension.

        Returns:
            `Shape`
        """
        assert self.rank == 1, "Shape.with_size() is only defined for shapes of rank 1."
        return self.with_sizes([size])

    def with_sizes(self, sizes: Union[Sequence[int], Sequence[Tuple[str, ...]], 'Shape', int], keep_item_names=True):
        """
        Returns a new `Shape` matching the dimension names and types of `self` but with different sizes.

        See Also:
            `Shape.with_size()`.

        Args:
            sizes: One of

                * `tuple` / `list` of same length as `self` containing replacement sizes or replacement item names.
                * `Shape` of any rank. Replaces sizes for dimensions shared by `sizes` and `self`.
                * `int`: new size for all dimensions

            keep_item_names: If `False`, forgets all item names.
                If `True`, keeps item names where the size does not change.

        Returns:
            `Shape` with same names and types as `self`.
        """
        if isinstance(sizes, (int, str)):
            sizes = [sizes] * len(self.sizes)
        if isinstance(sizes, Shape):
            item_names = [sizes.get_item_names(dim) if dim in sizes else self.get_item_names(dim) for dim in self.names]
            sizes = [sizes.get_size(dim) if dim in sizes else s for dim, s in self._named_sizes]
            return Shape(tuple(sizes), self.names, self.types, tuple(item_names))
        else:
            assert len(sizes) == len(self.sizes), f"Failed to set sizes of Shape {self} to {sizes} because rank does not match."
            sizes_ = []
            item_names = []
            for i, obj in enumerate(sizes):
                new_size, new_item_names = Shape._size_and_item_names_from_obj(obj, self.sizes[i], self.item_names[i], keep_item_names)
                sizes_.append(new_size)
                item_names.append(new_item_names)
            return Shape(tuple(sizes_), self.names, self.types, tuple(item_names))

    @staticmethod
    def _size_and_item_names_from_obj(obj, prev_size, prev_item_names, keep_item_names=True):
        if isinstance(obj, str):
            obj = [s.strip() for s in obj.split(',')]
        if isinstance(obj, (tuple, list)):
            return len(obj), tuple(obj)
        elif isinstance(obj, Number):
            return obj, prev_item_names if keep_item_names and (prev_size is None or _size_equal(obj, prev_size)) else None
        elif isinstance(obj, math.Tensor) or obj is None:
            return obj, None
        elif isinstance(obj, Shape):
            return obj.rank, obj.names
        else:
            raise ValueError(f"sizes can only contain int, str or Tensor but got {type(obj)}")

    def without_sizes(self):
        """
        Returns:
            `Shape` with all sizes undefined (`None`)
        """
        return Shape((None,) * self.rank, self.names, self.types, (None,) * self.rank)

    def _replace_single_size(self, dim: str, size: int, keep_item_names: bool = False):
        new_sizes = list(self.sizes)
        new_sizes[self.index(dim)] = size
        return self.with_sizes(new_sizes, keep_item_names=keep_item_names)

    def with_dim_size(self, dim: Union[str, 'Shape'], size: Union[int, 'math.Tensor', str, tuple, list], keep_item_names=True):
        """
        Returns a new `Shape` that has a different size for `dim`.

        Args:
            dim: Dimension for which to replace the size, `Shape` or `str`.
            size: New size, `int` or `Tensor`

        Returns:
            `Shape` with same names and types as `self`.
        """
        if isinstance(dim, Shape):
            dim = dim.name
        assert isinstance(dim, str)
        new_size, new_item_names = Shape._size_and_item_names_from_obj(size, self.get_size(dim), self.get_item_names(dim), keep_item_names)
        return self.replace(dim, Shape((new_size,), (dim,), (self.get_type(dim),), (new_item_names,)), keep_item_names=keep_item_names)

    def _with_names(self, names: Union[str, tuple, list]):
        if isinstance(names, str):
            names = parse_dim_names(names, self.rank)
            names = [n if n is not None else o for n, o in zip(names, self.names)]
        return Shape(self.sizes, tuple(names), self.types, self.item_names)

    def _replace_names_and_types(self,
                                 dims: Union['Shape', str, tuple, list],
                                 new: Union['Shape', str, tuple, list]) -> 'Shape':
        """
        Returns a copy of `self` with `dims` replaced by `new`.
        Dimensions that are not present in `self` are ignored.

        The dimension order is preserved.

        Args:
            dims: Dimensions to replace.
            new: New dimensions, must have same length as `dims`.
                If a `Shape` is given, replaces the dimension types and item names as well.

        Returns:
            `Shape` with same rank and dimension order as `self`.
        """
        dims = parse_dim_order(dims)
        sizes = [math.rename_dims(s, dims, new) if isinstance(s, math.Tensor) else s for s in self.sizes]
        new = parse_dim_order(new) if isinstance(new, str) else new
        names = list(self.names)
        types = list(self.types)
        item_names = list(self.item_names)
        for old_name, new_dim in zip(dims, new):
            if old_name in self:
                if isinstance(new_dim, Shape):
                    names[self.index(old_name)] = new_dim.name
                    types[self.index(old_name)] = new_dim.type
                    item_names[self.index(old_name)] = new_dim.item_names[0]
                else:
                    names[self.index(old_name)] = _apply_prefix(new_dim, types[self.index(old_name)])
        return Shape(tuple(sizes), tuple(names), tuple(types), tuple(item_names))

    def replace(self, dims: Union['Shape', str, tuple, list], new: 'Shape', keep_item_names=True, replace_item_names: DimFilter = None) -> 'Shape':
        """
        Returns a copy of `self` with `dims` replaced by `new`.
        Dimensions that are not present in `self` are ignored.

        The dimension order is preserved.

        Args:
            dims: Dimensions to replace.
            new: New dimensions, must have same length as `dims`.
                If a `Shape` is given, replaces the dimension types and item names as well.
            keep_item_names: Keeps existing item names for dimensions where `new` does not specify item names if the new dimension has the same size.
            replace_item_names: For which dims the item names should be replaced as well.

        Returns:
            `Shape` with same rank and dimension order as `self`.
        """
        dims = parse_dim_order(dims)
        assert isinstance(new, Shape), f"new must be a Shape but got {new}"
        names = list(self.names)
        sizes = list(self.sizes)
        types = list(self.types)
        item_names = list(self.item_names)
        for i in self.indices(self.only(replace_item_names)):
            if item_names[i]:
                if len(new) > len(dims):
                    raise NotImplementedError
                else:
                    name_map = {d: n for d, n in zip(dims, new.names)}
                    item_names[i] = tuple([name_map.get(n, n) for n in item_names[i]])
        if len(new) > len(dims):  # Put all in one spot
            assert len(dims) == 1, "Cannot replace 2+ dims by more replacements"
            index = self.index(dims[0])
            return concat_shapes(self[:index], new, self[index+1:])
        for old_name, new_dim in zip(dims, new):
            if old_name in self:
                names[self.index(old_name)] = new_dim.name
                types[self.index(old_name)] = new_dim.type
                if new_dim.item_names[0]:
                    item_names[self.index(old_name)] = new_dim.item_names[0]
                elif not _size_equal(new_dim.size, self.get_size(old_name)) or not keep_item_names:
                    item_names[self.index(old_name)] = None  # forget previous item names
                sizes[self.index(old_name)] = new_dim.size
        replaced = Shape(tuple(sizes), tuple(names), tuple(types), tuple(item_names))
        if len(new) == len(dims):
            return replaced
        to_remove = dims[-(len(dims) - len(new)):]
        return replaced.without(to_remove)

    def _with_types(self, types: Union['Shape', str, Tuple[str, ...], List[str]]):
        """
        Only for internal use.
        Note: This method does not rename dimensions to comply with type requirements (e.g. ~ for dual dims).
        """
        if isinstance(types, Shape):
            types = tuple([types.get_type(name) if name in types else self_type for name, self_type in zip(self.names, self.types)])
        elif isinstance(types, str):
            types = (types,) * self.rank
        elif isinstance(types, (tuple, list)):
            types = tuple(types)
        else:
            raise ValueError(types)
        names = tuple([_apply_prefix(name, t) for name, t in zip(self.names, types)])
        return Shape(self.sizes, names, types, self.item_names)

    def _with_item_names(self, item_names: tuple):
        return Shape(self.sizes, self.names, self.types, item_names)

    def _with_item_name(self, dim: str, item_name: tuple):
        if dim not in self:
            return self
        item_names = list(self.item_names)
        item_names[self.index(dim)] = item_name
        return Shape(self.sizes, self.names, self.types, tuple(item_names))

    def _perm(self, names: Tuple[str]) -> List[int]:
        assert len(set(names)) == len(names), f"No duplicates allowed but got {names}"
        assert len(names) >= len(self.names), f"Cannot find permutation for {self} given {names} because names {set(self.names) - set(names)} are missing"
        assert len(names) <= len(self.names), f"Cannot find permutation for {self} given {names} because too many names were passed: {names}"
        perm = [self.names.index(name) for name in names]
        return perm

    @property
    def volume(self) -> Union[int, None]:
        """
        Returns the total number of values contained in a tensor of this shape.
        This is the product of all dimension sizes.

        Returns:
            volume as `int` or `Tensor` or `None` if the shape is not `Shape.well_defined`
        """
        result = 1
        for size in self.sizes:
            if size is None:
                return None
            result *= size
        from ._tensors import Tensor
        if not isinstance(result, Tensor):
            return result
        result /= self.non_uniform_shape.volume  # We summed up the items -> undo multiplication
        return int(result.sum)

    @property
    def is_empty(self) -> bool:
        """ True if this shape has no dimensions. Equivalent to `Shape.rank` `== 0`. """
        return len(self.sizes) == 0

    def after_pad(self, widths: dict) -> 'Shape':
        sizes = list(self.sizes)
        item_names = list(self.item_names)
        for dim, (lo, up) in widths.items():
            if dim in self.names:
                sizes[self.index(dim)] += lo + up
                item_names[self.index(dim)] = None
        return Shape(tuple(sizes), self.names, self.types, tuple(item_names))

    def prepare_gather(self, dim: str, selection: Union[slice, int, 'Shape', str, tuple, list]) -> Union[slice, List[int]]:
        """
        Parse a slice object for a specific dimension.

        Args:
            dim: Name of dimension to slice.
            selection: Slice object.

        Returns:

        """
        if isinstance(selection, Shape):
            selection = selection.name if selection.rank == 1 else selection.names
        if isinstance(selection, str) and ',' in selection:
            selection = parse_dim_order(selection)
        if isinstance(selection, str):  # single item name
            item_names = self.get_item_names(dim, fallback_spatial=True)
            assert item_names is not None, f"No item names defined for dim '{dim}' in tensor {self.shape} and dimension size does not match spatial rank."
            assert selection in item_names, f"Accessing tensor.{dim}['{selection}'] failed. Item names are {item_names}."
            selection = item_names.index(selection)
        if isinstance(selection, (tuple, list)):
            selection = list(selection)
            if any([isinstance(s, str) for s in selection]):
                item_names = self.get_item_names(dim, fallback_spatial=True)
                for i, s in enumerate(selection):
                    if isinstance(s, str):
                        assert item_names is not None, f"Accessing tensor.{dim}['{s}'] failed because no item names are present on tensor {self.shape}"
                        assert s in item_names, f"Accessing tensor.{dim}['{s}'] failed. Item names are {item_names}."
                        selection[i] = item_names.index(s)
            if not selection:  # empty
                selection = slice(0, 0)
        return selection

    def prepare_renaming_gather(self, dim: str, selection: Union[slice, int, 'Shape', str, tuple, list]):
        if isinstance(selection, str) and '->' in selection:
            selection, new_names = selection.split('->')
            if new_names == '?':
                return self.prepare_gather(dim, selection), self[dim]._with_item_names((None,))
            else:
                return self.prepare_gather(dim, selection), self[dim].with_size(new_names)
        else:
            return self.prepare_gather(dim, selection), None

    def resolve_index(self, index: Dict[str, Union[slice, int, 'Shape', str, tuple, list]]) -> Dict[str, Union[slice, int, tuple, list]]:
        """
        Replaces item names by the corresponding indices.

        Args:
            index: n-dimensional index or slice.

        Returns:
            Same index but without any reference to item names.
        """
        return {dim: self.prepare_gather(dim, s) for dim, s in index.items()}

    def after_gather(self, selection: dict) -> 'Shape':
        from . import Tensor
        if self.is_non_uniform:
            sizes = [(s[selection] if isinstance(s, Tensor) else s) for s in self.sizes]
            sizes = [(int(s) if isinstance(s, Tensor) and s.rank == 0 else s) for s in sizes]
            result = self.with_sizes(sizes)
        else:
            result = self
        for sel_dim, sel in selection.items():
            if sel_dim not in self.names:
                continue
            sel = self.prepare_gather(sel_dim, sel)
            if isinstance(sel, int):
                result = result.without(sel_dim)
            elif isinstance(sel, slice):
                step = int(sel.step) if sel.step is not None else 1
                start = int(sel.start) if sel.start is not None else (0 if step > 0 else self.get_size(sel_dim)-1)
                stop = int(sel.stop) if sel.stop is not None else (self.get_size(sel_dim) if step > 0 else -1)
                if stop < 0 and step > 0:
                    stop += self.get_size(sel_dim)
                    assert stop >= 0
                if start < 0 and step > 0:
                    start += self.get_size(sel_dim)
                    assert start >= 0
                stop = min(stop, self.get_size(sel_dim))
                new_size = math.to_int64(math.ceil(math.wrap((stop - start) / step)))
                if new_size.rank == 0:
                    new_size = int(new_size)  # NumPy array not allowed because not hashable
                result = result._replace_single_size(sel_dim, new_size, keep_item_names=True)
                if step < 0:
                    result = result.flipped([sel_dim])
                if self.get_item_names(sel_dim) is not None:
                    result = result._with_item_name(sel_dim, tuple(self.get_item_names(sel_dim)[sel]))
            elif isinstance(sel, (tuple, list)):
                result = result._replace_single_size(sel_dim, len(sel))
                if self.get_item_names(sel_dim) is not None:
                    result = result._with_item_name(sel_dim, tuple([self.get_item_names(sel_dim)[i] for i in sel]))
            elif isinstance(sel, Tensor):
                if sel.dtype.kind == bool:
                    raise NotImplementedError("Shape.after_gather(Tensor[bool]) not yet implemented")
                    # from ._ops import nonzero
                    # sel = nonzero(sel)
                if sel.dtype.kind == int:
                    assert len(selection) == 1, f"When slicing a Shape with Tensor[int], only one sel item is allowed but got {sel}"
                    sel_shape = shape(sel)
                    assert sel_shape.channel_rank == 1 and sel_shape.channel.item_names[0], f"Shape.after_gather(Tensor[int]) requires indices to have a single channel dim with item names but got {sel}"
                    indexed = sel_shape.channel.item_names[0]
                    assert indexed in self, f"All indexed dims {indexed} must be part of sliced Shape {self}"
                    from ._ops import slice_
                    sizes = [slice_(s, sel) for s in self.sizes]
                    return self.with_sizes(sizes).without(indexed) & sel_shape.non_channel
            else:
                raise NotImplementedError(f"{type(sel)} not supported. Only (int, slice) allowed.")
        return result

    def meshgrid(self, names=False):
        """
        Builds a sequence containing all multi-indices within a tensor of this shape.
        All indices are returned as `dict` mapping dimension names to `int` indices.

        The corresponding values can be retrieved from Tensors and other Sliceables using `tensor[index]`.

        This function currently only supports uniform tensors.

        Args:
            names: If `True`, replace indices by their item names if available.

        Returns:
            `dict` iterator.
        """
        assert self.is_uniform, f"Shape.meshgrid() is currently not supported for non-uniform tensors, {self}"
        indices = [0] * self.rank
        while True:
            if names:
                yield {dim: (names[index] if names is not None else index) for dim, index, names in zip(self.names, indices, self.item_names)}
            else:
                yield {dim: index for dim, index in zip(self.names, indices)}
            for i in range(self.rank-1, -1, -1):
                indices[i] = (indices[i] + 1) % self.sizes[i]
                if indices[i] != 0:
                    break
            else:
                return

    def first_index(self, names=False):
        return next(iter(self.meshgrid(names=names)))

    def are_adjacent(self, dims: Union[str, tuple, list, set, 'Shape']):
        indices = self.indices(dims)
        return (max(indices) - min(indices)) == len(dims) - 1

    def __add__(self, other):
        if isinstance(other, Shape) and self.isdisjoint(other):
            return concat_shapes(self, other)
        return self._op2(other, lambda s, o: s + o, 0)

    def __radd__(self, other):
        return self._op2(other, lambda s, o: o + s, 0)

    def __sub__(self, other):
        if isinstance(other, (str, Shape, tuple, list, set)) or callable(other):
            return self.without(other)
        return self._op2(other, lambda s, o: s - o, 0)

    def __rsub__(self, other):
        return self._op2(other, lambda s, o: o - s, 0)

    def __mul__(self, other):
        return self._op2(other, lambda s, o: s * o, 1)

    def __rmul__(self, other):
        return self._op2(other, lambda s, o: o * s, 1)

    def _op2(self, other, fun, default: int):
        if isinstance(other, int):
            return Shape(tuple([fun(s, other) for s in self.sizes]), self.names, self.types, (None,) * self.rank)
        else:
            return NotImplemented

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

    @staticmethod
    def __stack__(values, dim: 'Shape', **kwargs):
        return shape_stack(dim, *values)


EMPTY_SHAPE = Shape((), (), (), ())
""" Empty shape, `()` """


class IncompatibleShapes(Exception):
    """
    Raised when the shape of a tensor does not match the other arguments.
    """
    def __init__(self, message, *shapes: Shape):
        Exception.__init__(self, message)
        self.shapes = shapes


def parse_dim_names(obj: Union[str, Sequence[str], Shape], count: int) -> tuple:
    if isinstance(obj, str):
        parts = obj.split(',')
        result = []
        for part in parts:
            part = part.strip()
            if part == '...':
                result.extend([None] * (count - len(parts) + 1))
            elif part == ':':
                result.append(None)
            else:
                result.append(part)
        assert len(result) == count, f"Number of specified names in '{obj}' does not match number of dimensions ({count})"
        return tuple(result)
    elif isinstance(obj, Shape):
        assert len(obj) == count, f"Number of specified names in {obj} does not match number of dimensions ({count})"
        return obj.names
    elif isinstance(obj, Sequence):
        assert len(obj) == count, f"Number of specified names in {obj} does not match number of dimensions ({count})"
        return tuple(obj)
    raise ValueError(obj)


def parse_dim_order(order: Union[str, tuple, list, Shape, None]) -> Union[tuple, None]:
    if isinstance(order, Shape):
        return order.names
    if isinstance(order, list):
        return tuple(order)
    elif isinstance(order, tuple):
        return order
    elif isinstance(order, str):
        parts = order.split(',')
        parts = [p.strip() for p in parts if p]
        return tuple(parts)
    raise ValueError(order)


def _construct_shape(dim_type: str, *args, **dims):
    sizes = ()
    names = []
    item_names = ()
    for arg in args:
        parts = [s.strip() for s in arg.split(',')]
        for name in parts:
            assert name not in names, f"Duplicate dimension name {name}"
            sizes += (None,)
            names.append(name)
            item_names += (None,)
    for name, size in dims.items():
        assert name not in names, f"Duplicate dimension name {name}"
        if isinstance(size, str):
            items = tuple([i.strip() for i in size.split(',')])
            if not items[-1]:
                items = items[:-1]
            size = len(items)
        elif isinstance(size, (tuple, list)):
            assert all(isinstance(s, str) for s in size), f"Item names must all be of type 'str' but got '{size}'"
            items = tuple(size)
            size = len(items)
        elif isinstance(size, Shape):
            items = size.names
            size = size.rank
        elif size is None or isinstance(size, int):
            # keep size
            items = None
        else:
            items = None
            from ._tensors import Tensor
            if isinstance(size, Tensor):
                size = int(size) if size.shape.volume == 1 else size
            else:
                try:
                    size = int(size)
                except ValueError:
                    raise ValueError(f"Cannot construct dimension from {type(size).__name__}. Only int, tuple, list, str or Shape allowed. Got {size}")
        names.append(name)
        sizes += (size,)
        item_names += (items,)
    names = tuple(_apply_prefix(name, dim_type) for name in names)
    return math.Shape(sizes, names, (dim_type,) * len(sizes), item_names)


def _apply_prefix(name: str, dim_type: str):
    match = re.search("\\w", name)
    assert match, f"Dimension name must contain at least one letter or underscore but got '{name}'"
    proper_name_index = match.start()
    prefix = '~' if dim_type == DUAL_DIM else ''
    return prefix + name[proper_name_index:]


def shape(obj, allow_unshaped=False) -> Shape:
    """
    If `obj` is a `Tensor` or `phiml.math.magic.Shaped`, returns its shape.
    If `obj` is a `Shape`, returns `obj`.

    This function can be passed as a `dim` argument to an operation to specify that it should act upon all dimensions.

    Args:
        obj: `Tensor` or `Shape` or `Shaped`
        allow_unshaped: If `True`, returns an empty shape for unsupported objects, else raises a `ValueError`.

    Returns:
        `Shape`
    """
    from .magic import PhiTreeNode, Shaped, BoundDim
    if isinstance(obj, Shape):
        return obj
    elif isinstance(obj, BoundDim):
        return shape(obj.obj)[obj.name]
    elif hasattr(obj, '__shape__'):
        return obj.__shape__()
    elif hasattr(obj, 'shape') and isinstance(obj.shape, Shape):
        return obj.shape
    elif isinstance(obj, (int, float, complex, bool)):
        return EMPTY_SHAPE
    elif isinstance(obj, (tuple, list)) and all(isinstance(item, (int, float, complex, bool)) for item in obj):
        return channel(vector=len(obj))
    elif isinstance(obj, (Number, bool)):
        return EMPTY_SHAPE
    elif obj is None:
        return EMPTY_SHAPE
    elif isinstance(obj, (tuple, list)) and all(isinstance(item, (PhiTreeNode, Shaped)) for item in obj):
        return merge_shapes(*obj, allow_varying_sizes=True)
    if isinstance(obj, dict) and all(isinstance(item, (PhiTreeNode, Shaped)) for item in obj):
        return merge_shapes(*obj.values(), allow_varying_sizes=True)
    elif isinstance(obj, PhiTreeNode):
        from ._magic_ops import all_attributes
        return merge_shapes(*[getattr(obj, a) for a in all_attributes(obj, assert_any=True)], allow_varying_sizes=True)
    else:
        from ..backend import choose_backend, NoBackendFound
        try:
            backend = choose_backend(obj)
            shape_tuple = backend.staticshape(obj)
            if len(shape_tuple) == 0:
                return EMPTY_SHAPE
            elif len(shape_tuple) == 1:
                return channel('vector')
            else:
                raise ValueError(f"Cannot auto-complete shape of {backend} tensor with shape {shape_tuple}. Only 0D and 1D tensors have a Φ-ML shape by default.")
        except NoBackendFound:
            if allow_unshaped:
                return EMPTY_SHAPE
            raise ValueError(f'shape() requires Shaped or Shape argument but got {type(obj)}')


def spatial(*args, **dims: Union[int, str, tuple, list, Shape, 'Tensor']) -> Shape:
    """
    Returns the spatial dimensions of an existing `Shape` or creates a new `Shape` with only spatial dimensions.

    Usage for filtering spatial dimensions:
    >>> spatial_dims = spatial(shape)
    >>> spatial_dims = spatial(tensor)

    Usage for creating a `Shape` with only spatial dimensions:
    >>> spatial_shape = spatial('undef', x=2, y=3)
    (x=2, y=3, undef=None)

    Here, the dimension `undef` is created with an undefined size of `None`.
    Undefined sizes are automatically filled in by `tensor`, `wrap`, `stack` and `concat`.

    To create a shape with multiple types, use `merge_shapes()`, `concat_shapes()` or the syntax `shape1 & shape2`.

    See Also:
        `channel`, `batch`, `instance`

    Args:
        *args: Either

            * `Shape` or `Tensor` to filter or
            * Names of dimensions with undefined sizes as `str`.

        **dims: Dimension sizes and names. Must be empty when used as a filter operation.

    Returns:
        `Shape` containing only dimensions of type spatial.
    """
    from .magic import Shaped
    if all(isinstance(arg, str) for arg in args) or dims:
        return _construct_shape(SPATIAL_DIM, *args, **dims)
    elif len(args) == 1 and isinstance(args[0], Shape):
        return args[0].spatial
    elif len(args) == 1 and isinstance(args[0], Shaped):
        return shape(args[0]).spatial
    else:
        raise AssertionError(f"spatial() must be called either as a selector spatial(Shape) or spatial(Tensor) or as a constructor spatial(*names, **dims). Got *args={args}, **dims={dims}")


def channel(*args, **dims: Union[int, str, tuple, list, Shape, 'Tensor']) -> Shape:
    """
    Returns the channel dimensions of an existing `Shape` or creates a new `Shape` with only channel dimensions.

    Usage for filtering channel dimensions:
    >>> channel_dims = channel(shape)
    >>> channel_dims = channel(tensor)

    Usage for creating a `Shape` with only channel dimensions:
    >>> channel_shape = channel('undef', vector=2)
    (vector=2, undef=None)

    Here, the dimension `undef` is created with an undefined size of `None`.
    Undefined sizes are automatically filled in by `tensor`, `wrap`, `stack` and `concat`.

    To create a shape with multiple types, use `merge_shapes()`, `concat_shapes()` or the syntax `shape1 & shape2`.

    See Also:
        `spatial`, `batch`, `instance`

    Args:
        *args: Either

            * `Shape` or `Tensor` to filter or
            * Names of dimensions with undefined sizes as `str`.

        **dims: Dimension sizes and names. Must be empty when used as a filter operation.

    Returns:
        `Shape` containing only dimensions of type channel.
    """
    from .magic import Shaped
    if all(isinstance(arg, str) for arg in args) or dims:
        return _construct_shape(CHANNEL_DIM, *args, **dims)
    elif len(args) == 1 and isinstance(args[0], Shape):
        return args[0].channel
    elif len(args) == 1 and isinstance(args[0], Shaped):
        return shape(args[0]).channel
    else:
        raise AssertionError(f"channel() must be called either as a selector channel(Shape) or channel(Tensor) or as a constructor channel(*names, **dims). Got *args={args}, **dims={dims}")


def batch(*args, **dims: Union[int, str, tuple, list, Shape, 'Tensor']) -> Shape:
    """
    Returns the batch dimensions of an existing `Shape` or creates a new `Shape` with only batch dimensions.

    Usage for filtering batch dimensions:
    >>> batch_dims = batch(shape)
    >>> batch_dims = batch(tensor)

    Usage for creating a `Shape` with only batch dimensions:
    >>> batch_shape = batch('undef', batch=2)
    (batch=2, undef=None)

    Here, the dimension `undef` is created with an undefined size of `None`.
    Undefined sizes are automatically filled in by `tensor`, `wrap`, `stack` and `concat`.

    To create a shape with multiple types, use `merge_shapes()`, `concat_shapes()` or the syntax `shape1 & shape2`.

    See Also:
        `channel`, `spatial`, `instance`

    Args:
        *args: Either

            * `Shape` or `Tensor` to filter or
            * Names of dimensions with undefined sizes as `str`.

        **dims: Dimension sizes and names. Must be empty when used as a filter operation.

    Returns:
        `Shape` containing only dimensions of type batch.
    """
    from .magic import Shaped
    if all(isinstance(arg, str) for arg in args) or dims:
        return _construct_shape(BATCH_DIM, *args, **dims)
    elif len(args) == 1 and isinstance(args[0], Shape):
        return args[0].batch
    elif len(args) == 1 and isinstance(args[0], Shaped):
        return shape(args[0]).batch
    else:
        raise AssertionError(f"batch() must be called either as a selector batch(Shape) or batch(Tensor) or as a constructor batch(*names, **dims). Got *args={args}, **dims={dims}")


def instance(*args, **dims: Union[int, str, tuple, list, Shape, 'Tensor']) -> Shape:
    """
    Returns the instance dimensions of an existing `Shape` or creates a new `Shape` with only instance dimensions.

    Usage for filtering instance dimensions:
    >>> instance_dims = instance(shape)
    >>> instance_dims = instance(tensor)

    Usage for creating a `Shape` with only instance dimensions:
    >>> instance_shape = instance('undef', points=2)
    (points=2, undef=None)

    Here, the dimension `undef` is created with an undefined size of `None`.
    Undefined sizes are automatically filled in by `tensor`, `wrap`, `stack` and `concat`.

    To create a shape with multiple types, use `merge_shapes()`, `concat_shapes()` or the syntax `shape1 & shape2`.

    See Also:
        `channel`, `batch`, `spatial`

    Args:
        *args: Either

            * `Shape` or `Tensor` to filter or
            * Names of dimensions with undefined sizes as `str`.

        **dims: Dimension sizes and names. Must be empty when used as a filter operation.

    Returns:
        `Shape` containing only dimensions of type instance.
    """
    from .magic import Shaped
    if all(isinstance(arg, str) for arg in args) or dims:
        return _construct_shape(INSTANCE_DIM, *args, **dims)
    elif len(args) == 1 and isinstance(args[0], Shape):
        return args[0].instance
    elif len(args) == 1 and isinstance(args[0], Shaped):
        return shape(args[0]).instance
    else:
        raise AssertionError(f"instance() must be called either as a selector instance(Shape) or instance(Tensor) or as a constructor instance(*names, **dims). Got *args={args}, **dims={dims}")


def dual(*args, **dims: Union[int, str, tuple, list, Shape, 'Tensor']) -> Shape:
    """
    Returns the dual dimensions of an existing `Shape` or creates a new `Shape` with only dual dimensions.

    Dual dimensions are assigned the prefix `~` to distinguish them from regular dimensions.
    This way, a regular and dual dimension of the same name can exist in one `Shape`.

    Dual dimensions represent the input space and are typically only present on matrices or higher-order matrices.
    Dual dimensions behave like batch dimensions in regular operations, if supported.
    During matrix multiplication, they are matched against their regular counterparts by name (ignoring the `~` prefix).

    Usage for filtering dual dimensions:

    >>> dual_dims = dual(shape)
    >>> dual_dims = dual(tensor)

    Usage for creating a `Shape` with only dual dimensions:

    >>> dual('undef', points=2)
    (~undefᵈ=None, ~pointsᵈ=2)

    Here, the dimension `undef` is created with an undefined size of `None`.
    Undefined sizes are automatically filled in by `tensor`, `wrap`, `stack` and `concat`.

    To create a shape with multiple types, use `merge_shapes()`, `concat_shapes()` or the syntax `shape1 & shape2`.

    See Also:
        `channel`, `batch`, `spatial`

    Args:
        *args: Either

            * `Shape` or `Tensor` to filter or
            * Names of dimensions with undefined sizes as `str`.

        **dims: Dimension sizes and names. Must be empty when used as a filter operation.

    Returns:
        `Shape` containing only dimensions of type dual.
    """
    from .magic import Shaped
    if all(isinstance(arg, str) for arg in args) or dims:
        return _construct_shape(DUAL_DIM, *args, **dims)
    elif len(args) == 1 and isinstance(args[0], Shape):
        return args[0].dual
    elif len(args) == 1 and isinstance(args[0], Shaped):
        return shape(args[0]).dual
    else:
        raise AssertionError(f"dual() must be called either as a selector dual(Shape) or dual(Tensor) or as a constructor dual(*names, **dims). Got *args={args}, **dims={dims}")
    

def auto(spec: Union[str, Shape], default_type: Callable = None) -> Shape:
    """
    Create a single-dimension `Shape` from a specification string.
    
    Args:
        spec: String specifying the dimension name and type.
            The type can be specified by a trailing superscript letter or `:` followed by the regular letter.
            Examples: `vector:c` or `vectorᶜ` indicate channel dimensions.
            Leading `~` indicate dual dimensions.
        default_type: Fallback type if no type is specified.
            If not provided, an error will be thrown if `spec` does not specify the type.

    Returns:
        `Shape`
    """
    if isinstance(spec, Shape):
        return spec  # allow multi-dim Shapes as well, as the main application is stacking
    assert isinstance(spec, str), f"spec must be a Shape or str but got {type(spec)}"
    assert ',' not in spec, f"auto dim only supported for single dimensions"
    return parse_shape_spec(spec, default_type=default_type)


class InvalidShapeSpec(ValueError):
    pass


SPEC_PATTERNS = {
    'name_type_items': re.compile(r'(~?)(\w+):(\w+)=\(([^)]*)\)'),
    'name_type': re.compile(r'(~?)(\w+):(\w+)'),
    'name_items': re.compile(r'(~?)(\w+)=\(([^)]*)\)'),
    'items': re.compile(r'(~?)\(([^)]*)\)'),
    'dual_name': re.compile(r'(~\w+)'),
    'single_letter': re.compile(r'(\w)(?=,|$)'),
    'name_only': re.compile(r'(\w+)(?=,|$)')
}


def parse_shape_spec(input_string, default_type: Callable = None) -> Shape:
    results = []
    pos = 0
    while pos < len(input_string):
        if match := SPEC_PATTERNS['name_type_items'].match(input_string, pos):
            tilde, name, type_, values = match.groups()
            if tilde and type_ not in ('d', 'dual'):
                raise InvalidShapeSpec(input_string, f"Dimension names starting with ~ must be of type dual. Failed at index {pos}: {input_string[pos:]}")
            elif not tilde and type_ in ('d', 'dual'):
                raise InvalidShapeSpec(input_string, f"Dual dims must start with ~. Failed at index {pos}: {input_string[pos:]}")
            items = [n.strip() for n in values.split(',') if n.strip()]
            results.append({'name': '~' + name if tilde else name, 'type': type_, 'values': items})
            pos = match.end() + 1
        elif match := SPEC_PATTERNS['name_type'].match(input_string, pos):
            tilde, name, type_ = match.groups()
            if tilde and type_ not in ('d', 'dual'):
                raise InvalidShapeSpec(input_string, f"Dimension names starting with ~ must be of type dual. Failed at index {pos}: {input_string[pos:]}")
            elif not tilde and type_ in ('d', 'dual'):
                raise InvalidShapeSpec(input_string, f"Dual dims must start with ~. Failed at index {pos}: {input_string[pos:]}")
            # Check if the next character is an equal sign followed by parentheses
            next_char_pos = pos + len(match.group())
            if next_char_pos < len(input_string) and input_string[next_char_pos] == '=':
                raise ValueError(f"Invalid format at position {pos}: values must be inside parentheses")
            results.append({'name': '~' + name if tilde else name, 'type': type_})
            pos = match.end() + 1
        elif match := SPEC_PATTERNS['name_items'].match(input_string, pos):
            tilde, name, values = match.groups()
            items = [n.strip() for n in values.split(',') if n.strip()]
            results.append({'name': '~' + name if tilde else name, 'type': 'd' if tilde else 'c', 'values': items})
            pos = match.end() + 1
        elif match := SPEC_PATTERNS['items'].match(input_string, pos):
            tilde, values = match.groups()
            results.append({'name': '~vector' if tilde else 'vector', 'type': 'd' if tilde else 'c', 'values': values.split(',')})
            pos = match.end() + 1
        elif match := SPEC_PATTERNS['dual_name'].match(input_string, pos):
            name, = match.groups()
            results.append({'name': name, 'type': 'd'})
            pos = match.end() + 1
        elif match := SPEC_PATTERNS['single_letter'].match(input_string, pos):
            name, = match.groups()
            results.append({'name': name, 'type': 's'})
            pos = match.end() + 1
        elif default_type is not None and (match := SPEC_PATTERNS['name_only'].match(input_string, pos)):
            name, = match.groups()
            default_type_str = TYPE_BY_FUNCTION[default_type]
            results.append({'name': name, 'type': default_type_str})
            pos = match.end() + 1
        else:
            raise InvalidShapeSpec(input_string, f"Failed to parse from index {pos}: '{input_string[pos:]}'. Dims must be specified as name:type or name:type=(item_names...). Names and types may only be omitted if component names are given.")
    names = [r['name'] for r in results]
    types = [r['type'] for r in results]
    types = [INV_CHAR[t] if len(t) == 1 else t for t in types]
    item_names = [r.get('values', None) for r in results]
    item_names = [tuple(items) if items is not None else None for items in item_names]
    sizes = [len(items) if items is not None else None for items in item_names]
    return Shape(tuple(sizes), tuple(names), tuple(types), tuple(item_names))


DIM_FUNCTIONS = {BATCH_DIM: batch, SPATIAL_DIM: spatial, INSTANCE_DIM: instance, CHANNEL_DIM: channel, DUAL_DIM: dual}
TYPE_BY_FUNCTION = {v: k for k, v in DIM_FUNCTIONS.items()}


def merge_shapes(*objs: Union[Shape, Any], order=(batch, dual, instance, spatial, channel), allow_varying_sizes=False):
    """
    Combines `shapes` into a single `Shape`, grouping dimensions by type.
    If dimensions with equal names are present in multiple shapes, their types and sizes must match.

    The shorthand `shape1 & shape2` merges shapes with `check_exact=[spatial]`.

    See Also:
        `concat_shapes()`.

    Args:
        *objs: `Shape` or `Shaped` objects to combine.
        order: Dimension type order as `tuple` of type filters (`channel`, `batch`, `spatial` or `instance`). Dimensions are grouped by type while merging.

    Returns:
        Merged `Shape`

    Raises:
        IncompatibleShapes if the shapes are not compatible
    """
    if not objs:
        return EMPTY_SHAPE
    shapes = [obj if isinstance(obj, Shape) else shape(obj) for obj in objs]
    merged = []
    for dim_type in order:
        type_group = dim_type(shapes[0])
        for sh in shapes[1:]:
            sh = dim_type(sh)
            for dim in sh:
                if dim not in type_group:
                    type_group = type_group._expand(dim, pos=-1)
                else:  # check size match
                    sizes_match = _size_equal(dim.size, type_group.get_size(dim.name))
                    if allow_varying_sizes:
                        if not sizes_match:
                            type_group = type_group.with_dim_size(dim, None)
                    else:
                        if not sizes_match:
                            raise IncompatibleShapes(f"Cannot merge shapes {shapes} because dimension '{dim.name}' exists with different sizes.", *shapes)
                        names1 = type_group.get_item_names(dim)
                        names2 = sh.get_item_names(dim)
                        if names1 is not None and names2 is not None and len(names1) > 1:
                            if names1 != names2:
                                if set(names1) == set(names2):
                                    raise IncompatibleShapes(f"Inconsistent component order on {dim.name}: '{','.join(names1)}' vs '{','.join(names2)}' in dimension '{dim.name}'. Failed to merge shapes {shapes}", *shapes)
                                else:
                                    raise IncompatibleShapes(f"Cannot merge shapes {shapes} because dimension '{dim.name}' exists with different item names.", *shapes)
                        elif names1 is None and names2 is not None:
                            type_group = type_group._with_item_name(dim, tuple(names2))
        merged.append(type_group)
    return concat_shapes(*merged)


def non_batch(obj) -> Shape:
    """
    Returns the non-batch dimensions of an object.

    Args:
        obj: `Shape` or object with a valid `shape` property.

    Returns:
        `Shape`
    """
    from .magic import Shaped
    if isinstance(obj, Shape):
        return obj.non_batch
    elif isinstance(obj, Shaped):
        return shape(obj).non_batch
    else:
        raise AssertionError(f"non_batch() must be called either on a Shape or an object with a 'shape' property but got {obj}")


def non_spatial(obj) -> Shape:
    """
    Returns the non-spatial dimensions of an object.

    Args:
        obj: `Shape` or object with a valid `shape` property.

    Returns:
        `Shape`
    """
    from .magic import Shaped
    if isinstance(obj, Shape):
        return obj.non_spatial
    elif isinstance(obj, Shaped):
        return shape(obj).non_spatial
    else:
        raise AssertionError(f"non_spatial() must be called either on a Shape or an object with a 'shape' property but got {obj}")


def non_instance(obj) -> Shape:
    """
    Returns the non-instance dimensions of an object.

    Args:
        obj: `Shape` or object with a valid `shape` property.

    Returns:
        `Shape`
    """
    from .magic import Shaped
    if isinstance(obj, Shape):
        return obj.non_instance
    elif isinstance(obj, Shaped):
        return shape(obj).non_instance
    else:
        raise AssertionError(f"non_instance() must be called either on a Shape or an object with a 'shape' property but got {obj}")


def non_channel(obj) -> Shape:
    """
    Returns the non-channel dimensions of an object.

    Args:
        obj: `Shape` or object with a valid `shape` property.

    Returns:
        `Shape`
    """
    from .magic import Shaped
    if isinstance(obj, Shape):
        return obj.non_channel
    elif isinstance(obj, Shaped):
        return shape(obj).non_channel
    else:
        raise AssertionError(f"non_channel() must be called either on a Shape or an object with a 'shape' property but got {obj}")


def non_dual(obj) -> Shape:
    """
    Returns the non-dual dimensions of an object.

    Args:
        obj: `Shape` or object with a valid `shape` property.

    Returns:
        `Shape`
    """
    from .magic import Shaped
    if isinstance(obj, Shape):
        return obj.non_dual
    elif isinstance(obj, Shaped):
        return shape(obj).non_dual
    else:
        raise AssertionError(f"non_dual() must be called either on a Shape or an object with a 'shape' property but got {obj}")


def non_primal(obj) -> Shape:
    """
    Returns the batch and dual dimensions of an object.

    Args:
        obj: `Shape` or object with a valid `shape` property.

    Returns:
        `Shape`
    """
    from .magic import Shaped
    if isinstance(obj, Shape):
        return obj.non_primal
    elif isinstance(obj, Shaped):
        return shape(obj).non_primal
    else:
        raise AssertionError(f"non_dual() must be called either on a Shape or an object with a 'shape' property but got {obj}")


def primal(obj) -> Shape:
    """
    Returns the instance, spatial and channel dimensions of an object.

    Args:
        obj: `Shape` or object with a valid `shape` property.

    Returns:
        `Shape`
    """
    from .magic import Shaped
    if isinstance(obj, Shape):
        return obj.primal
    elif isinstance(obj, Shaped):
        return shape(obj).primal
    else:
        raise AssertionError(f"primal() must be called either on a Shape or an object with a 'shape' property but got {obj}")


def _size_equal(s1, s2):
    if s1 is None:
        return s2 is None
    if s2 is None:
        return False
    if isinstance(s1, int):
        return isinstance(s2, int) and s2 == s1
    else:
        return math.close(s1, s2)


def concat_shapes(*shapes: Union[Shape, Any]) -> Shape:
    """
    Creates a `Shape` listing the dimensions of all `shapes` in the given order.

    See Also:
        `merge_shapes()`.

    Args:
        *shapes: Shapes to concatenate. No two shapes must contain a dimension with the same name.

    Returns:
        Combined `Shape`.
    """
    shapes = [obj if isinstance(obj, Shape) else shape(obj) for obj in shapes]
    names = sum([s.names for s in shapes], ())
    if len(set(names)) != len(names):
        raise IncompatibleShapes(f"Cannot concatenate shapes {list(shapes)}. Duplicate dimension names are not allowed.")
    sizes = sum([s.sizes for s in shapes], ())
    types = sum([s.types for s in shapes], ())
    item_names = sum([s.item_names for s in shapes], ())
    return Shape(sizes, names, types, item_names)


def shape_stack(stack_dim: Shape, *shapes: Shape, stack_dim_first=False):
    """ Returns the shape of a tensor created by stacking tensors with `shapes`. """
    if stack_dim.rank > 1:
        assert stack_dim.volume == len(shapes), f"stack_dim {stack_dim} does not match number of shapes: {len(shapes)}"
    names = list(stack_dim.names)
    types = list(stack_dim.types)
    item_names = list(stack_dim.item_names)
    incompatible_item_names = []
    for other in shapes:
        for size, name, type, items in other._dimensions:
            if name not in names:
                if type in types:
                    index = len(types) - types[::-1].index(type)
                elif type == BATCH_DIM:
                    index = 0
                elif type == DUAL_DIM:
                    index = min([len(names), *[i for i in range(len(names)) if types[i] == DUAL_DIM]])
                elif type == CHANNEL_DIM:
                    index = len(names)
                elif type == SPATIAL_DIM:
                    index = min([len(names), *[i for i in range(len(names)) if types[i] == CHANNEL_DIM]])
                elif type == INSTANCE_DIM:
                    index = min([len(names), *[i for i in range(len(names)) if types[i] == INSTANCE_DIM]])
                else:
                    raise ValueError(type)
                if stack_dim_first:
                    index = max(index, stack_dim.rank)
                names.insert(index, name)
                types.insert(index, type)
                item_names.insert(index, items)
            else:
                index = names.index(name)
                if items != item_names[index]:
                    if item_names[index] is None:
                        item_names[index] = items
                    else:
                        warnings.warn(f"Stacking shapes with incompatible item names will result in item names being lost. For {name} Got {item_names[index]} and {items}", RuntimeWarning)
                        incompatible_item_names.append(index)
    if incompatible_item_names:
        for index in incompatible_item_names:
            item_names[index] = None
    sizes = []
    for name in names:
        if name in stack_dim.names:
            if stack_dim.rank == 1:
                size = len(shapes)
            else:
                size = stack_dim.get_size(name)
        else:
            dim_sizes = [(s.get_size(name) if name in s else None) for s in shapes]
            valid_dim_sizes = [s.get_size(name) for s in shapes if name in s]
            if all([math.close(s, valid_dim_sizes[0]) for s in valid_dim_sizes[1:]]):
                size = valid_dim_sizes[0]
            else:
                from ._magic_ops import stack
                dim_sizes_or_1 = [1 if s is None else s for s in dim_sizes]
                size = stack(dim_sizes_or_1, stack_dim, expand_values=True)
        sizes.append(size)
    return Shape(tuple(sizes), tuple(names), tuple(types), tuple(item_names))


def vector_add(*shapes: Shape):
    if not shapes:
        return EMPTY_SHAPE
    names = shapes[0].names
    types = shapes[0].types
    item_names = shapes[0].item_names
    for shape in shapes[1:]:
        for name in shape.names:
            if name not in names:
                names += (name,)
                types += (shape.get_type(name),)
                item_names += (shape.get_item_names(name),)
    sizes = [sum(sh.get_size(dim) if dim in sh else 0 for sh in shapes) for dim in names]
    return Shape(tuple(sizes), names, types, item_names)
