Module phiml.math.magic

Magic methods allow custom classes to be compatible with various functions defined in phiml.math, analogous to how implementing __hash__ allows objects to be used with hash(). The magic methods are grouped into purely declarative classes (interfaces) by what functionality they provide.

• Shaped objects have a Shape.
• Sliceable objects can be sliced along dimensions.
• Shapable objects can additionally be reshaped.
• PhiTreeNode objects can be disassembled into tensors.

All of these magic classes declared here define a custom instance checks and should not be used as superclasses.

An object implements one of the types defined here by implementing one or more of the related magic methods. Instance checks can be performed via isinstance(obj, <MagicClass>).

This is analogous to interfaces defined in the built-in collections package, such as Sized, Iterable, Hashable, Callable. To check whether len(obj) can be performed, you check isinstance(obj, Sized).

"""
Magic methods allow custom classes to be compatible with various functions defined in `phiml.math`, analogous to how implementing `__hash__` allows objects to be used with `hash()`.
The magic methods are grouped into purely declarative classes (interfaces) by what functionality they provide.

* `Shaped` objects have a `phiml.math.Shape`.
* `Sliceable` objects can be sliced along dimensions.
* `Shapable` objects can additionally be reshaped.
* `PhiTreeNode` objects can be disassembled into tensors.

All of these magic classes declared here define a custom instance checks and should not be used as superclasses.

An object implements one of the types defined here by implementing one or more of the related magic methods.
Instance checks can be performed via `isinstance(obj, <MagicClass>)`.

This is analogous to interfaces defined in the built-in `collections` package, such as `Sized, Iterable, Hashable, Callable`.
To check whether `len(obj)` can be performed, you check `isinstance(obj, Sized)`.
"""
import warnings
from typing import Tuple, Callable, Union

import dataclasses

from ._shape import Shape, shape, channel, non_batch, batch, spatial, instance, concat_shapes, dual
from ..backend._dtype import DType


class _ShapedType(type):
    def __instancecheck__(self, instance):
        try:
            shape(instance)
            return True
        except ValueError:
            return False

    def __subclasscheck__(self, subclass):
        return True


class Shaped(metaclass=_ShapedType):
    """
    To be considered shaped, an object must either implement the magic method `__shape__()` or have a valid `shape` property.
    In either case, the returned shape must be an instance of `phiml.math.Shape`.

    To check whether an object is `Shaped`, use `isinstance(obj, Shaped)`.

    **Usage in `phiml.math`:**

    The functions `phiml.math.shape` as well as dimension filters, such as `phiml.math.spatial` or `phiml.math.non_batch` can be called on all shaped objects.

    See Also:
        `Sliceable`, `Shapable`
    """

    def __shape__(self) -> 'Shape':
        """
        Returns the shape of this object.

        Alternatively, the shape can be declared via the property `shape`.

        Returns:
            `phiml.math.Shape`
        """
        raise NotImplementedError

    @property
    def shape(self) -> 'Shape':
        """
        Alternative form of `__shape__()`.
        Implement either to be considered `Shaped`.

        Returns:
            `phiml.math.Shape`
        """
        raise NotImplementedError


class _SliceableType(type):
    def __instancecheck__(self, instance):
        if isinstance(instance, str):
            return False
        return isinstance(instance, Shaped) and (hasattr(instance, '__getitem__') or isinstance(instance, PhiTreeNode))

    def __subclasscheck__(self, subclass):
        if subclass == str:
            return False
        return hasattr(subclass, '__getitem__')


class Sliceable(metaclass=_SliceableType):
    """
    Objects are considered sliceable if they are `Shaped` and implement `__getitem__` as defined below.

    To enable the slicing syntax `obj.dim[slice]`, implement the `__getattr__` method as defined below.

    Classes implementing `Sliceable` should override `__getattr__` to enable the special slicing syntax defined in `BoundDim`.

    **Usage in `phiml.math`:**

    In addition to slicing, sliceable objects can be unstacked along one or multiple dimensions using `phiml.math.unstack`.

    See Also
        `Shapable`, `Shaped`
    """

    def __getitem__(self, item) -> 'Sliceable':
        """
        Slice this object along one or multiple existing or non-existing dimensions.

        When overriding this function, make sure to first call `slicing_dict(self, item)` to sort slices by dimension.

        Args:
            item: `dict` mapping dimension names to the corresponding selections.
                Selections can be slices, indices, tuples, item names, bool tensors, int tensors or other custom types.
                All Sliceable object must support indexing by `int`, `slice`, `tuple`, `list`, `str`.

        Returns:
            Instance of the same class (or a compatible class) as `self`.
        """
        raise NotImplementedError

    def __unstack__(self, dims: Tuple[str, ...]) -> Tuple['Sliceable', ...]:
        """
        Un-stack this object along one or multiple dimensions.
        Un-stacking along multiple dimensions is equal to first packing the dimensions and then unstacking along the packed dimension.

        Implementing this magic method is optional but the default implementation may be slow.

        Args:
            dims: Ordered `tuple` of dimension names along which to unstack this object.

        Returns:
            `tuple` of slices along `dims` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError


class _ShapableType(type):
    def __instancecheck__(self, instance):
        return isinstance(instance, Sliceable) and isinstance(instance, Shaped) and\
               (hasattr(instance, '__stack__') or (hasattr(instance, '__concat__') and hasattr(instance, '__expand__')) or isinstance(instance, PhiTreeNode))

    def __subclasscheck__(self, subclass):
        return issubclass(subclass, Sliceable) and\
               (hasattr(subclass, '__stack__') or (hasattr(subclass, '__concat__') and hasattr(subclass, '__expand__')) or issubclass(subclass, PhiTreeNode))


class Shapable(metaclass=_ShapableType):
    """
    Shapable objects can be stacked, concatenated and reshaped.

    To be considered `Shapable`, objects must be `Sliceable` and `Shaped` and implement

    * `__stack__()` or
    * `__concat__()` and `__expand__()`.

    Objects should additionally implement the other magic methods for performance reasons.

    **Usage in `phiml.math`:**

    Shapable objects can be used with the following functions in addition to what they inherit from being `Sliceable` and `Shaped`:

    * `phiml.math.stack`
    * `phiml.math.concat`
    * `phiml.math.expand`
    * `phiml.math.rename_dims`
    * `phiml.math.pack_dims`
    * `phiml.math.unpack_dim`
    * `phiml.math.flatten`

    Additionally, the `phiml.math.BoundDim` syntax for dimension renaming and retyping is enabled, e.g. `obj.dim.as_channel('vector')`.
    """
    @staticmethod
    def __stack__(values: tuple, dim: Shape, **kwargs) -> 'Shapable':
        """
        Stack all `values` into a single instance along the new dimension `dim`.

        This method can be implemented as a bound method or as a `staticmethod` (without the `self` argument).

        Args:
            values: `tuple` of `Shapable` objects to be stacked. `self` is included in that list at least once.
            dim: Single-dimension `Shape`. This dimension must not be present with any of the `values`.
                The dimension fulfills the condition `dim.size == len(values)`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` representing the stacked slices.
            Its shape includes `dim` in addition to the dimensions present in `values`.
            If such a representation cannot be created because some values in `values` are not supported, returns `NotImplemented`.
        """
        raise NotImplementedError

    @staticmethod
    def __concat__(values: tuple, dim: str, **kwargs) -> 'Shapable':
        """
        Concatenate `values` along `dim`.

        This method can be implemented as a bound method or as a `staticmethod` (without the `self` argument).

        Args:
            values: Values to concatenate. `self` is included in that list at least once.
            dim: Dimension nams as `str`, must be present in all `values`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` representing the concatenated values or `NotImplemented` to revert to default behavior for this object.
            When returning a valid object, the size of `dim` must be equal to the sum of all `dim` sizes in `values`.
            If such a representation cannot be created because some values in `values` are not supported, returns `NotImplemented`.
        """
        raise NotImplementedError

    def __expand__(self, dims: Shape, **kwargs) -> 'Shapable':
        """
        Adds new dimensions to this object.
        The value of this object is constant along the new dimensions.

        Args:
            dims: Dimensions to add.
                They are guaranteed to not already be present in `shape(self)`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

    def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Shapable':
        """
        Exchange existing dimensions.
        This can be used to rename dimensions, change dimension types or change item names.

        Args:
            dims: Dimensions to be replaced.
            new_dims: Replacement dimensions as `Shape` with `rank == len(dims)`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

    def __pack_dims__(self, dims: Tuple[str, ...], packed_dim: Shape, pos: Union[int, None], **kwargs) -> 'Shapable':
        """
        Compresses multiple dimensions into a single dimension by concatenating the elements.
        Elements along the new dimensions are laid out according to the order of `dims`.

        The type of the new dimension will be equal to the types of `dims`.
        If `dims` have varying types, the new dimension will be a batch dimension.

        Args:
            dims: Dimensions to be compressed in the specified order.
            packed_dim: Single-dimension `Shape`.
            pos: Index of new dimension. `None` for automatic, `-1` for last, `0` for first.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

    def __unpack_dim__(self, dim: str, unpacked_dims: Shape, **kwargs) -> 'Shapable':
        """
        Decompresses a tensor dimension by unstacking the elements along it.
        The compressed dimension `dim` is assumed to contain elements laid out according to the order of `unpacked_dims`.

        Args:
            dim: Dimension to be decompressed.
            unpacked_dims: `Shape`: Ordered dimensions to replace `dim`, fulfilling `unpacked_dims.volume == shape(self)[dim].rank`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

    def __flatten__(self, flat_dim: Shape, flatten_batch: bool, **kwargs) -> 'Shapable':
        """
        Lays out all elements along a single dimension.
        This is equivalent to packing all dimensions.

        Args:
            flat_dim: Single dimension as `Shape`.
            flatten_batch: Whether to flatten batch dimensions as well.
            If `False`, batch dimensions are kept, only onn-batch dimensions are flattened.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError


class _PhiTreeNodeType(type):

    def __instancecheck__(self, instance):
        from ._tensors import Tensor, MISSING_TENSOR, NATIVE_TENSOR, Dict
        if isinstance(instance, Tensor):
            return True
        if instance is MISSING_TENSOR or instance is NATIVE_TENSOR:
            return True
        if instance is None or isinstance(instance, Tensor):
            return True
        elif isinstance(instance, (tuple, list)):
            return all(isinstance(item, PhiTreeNode) for item in instance)
        elif isinstance(instance, Dict):
            return True
        elif isinstance(instance, dict):
            return all(isinstance(val, (Shaped, PhiTreeNode)) for val in instance.values())
        elif dataclasses.is_dataclass(instance):
            return True
        else:
            return hasattr(instance, '__variable_attrs__') or hasattr(instance, '__value_attrs__')

    def __subclasscheck__(self, subclass):
        from ._tensors import Tensor, Dict
        if issubclass(subclass, Tensor):
            return True
        if subclass in (tuple, list, dict):
            return True
        elif issubclass(subclass, Dict):
            return True
        elif dataclasses.is_dataclass(subclass):
            return True
        else:
            return hasattr(subclass, '__variable_attrs__') or hasattr(subclass, '__value_attrs__')


class PhiTreeNode(metaclass=_PhiTreeNodeType):
    """
    Φ-tree nodes can be iterated over and disassembled or flattened into elementary objects, such as tensors.
    `phiml.math.Tensor` instances as well as PyTree nodes (`tuple`, `list`, `dict` with `str` keys) are Φ-tree nodes.
    All data classes are also considered PhiTreeNodes as of version 2.3.

    For custom classes to be considered Φ-tree nodes, they have to be a dataclass or implement one of the following magic methods:

    * `__variable_attrs__()`
    * `__value_attrs__()`

    Dataclasses may also implement these functions to specify which attributes should be considered value / variable properties.

    Additionally, Φ-tree nodes must override `__eq__()` to allow comparison of data-stripped (key) instances.

    To check whether an object is a Φ-tree node, use `isinstance(obj, PhiTreeNode)`.

    **Usage in `phiml.math`:**

    Φ-tree nodes can be used as keys, for example in `jit_compile()`.
    They are converted to keys by stripping all variable tensors and replacing them by a placeholder object.
    In key mode, `__eq__()` compares all non-variable properties that might invalidate a trace when changed.

    Disassembly and assembly of Φ-tree nodes uses `phiml.math.copy_with` which will call `__with_attrs__` if implemented.
    """

    def __value_attrs__(self) -> Tuple[str, ...]:
        """
        Returns all `Tensor` or `PhiTreeNode` attribute names of `self` that should be transformed by single-operand math operations,
        such as `sin()`, `exp()`.

        Returns:
            `tuple` of `str` attributes.
                Calling `getattr(self, attr)` must return a `Tensor` or `PhiTreeNode` for all returned attributes.
        """
        raise NotImplementedError

    def __variable_attrs__(self) -> Tuple[str, ...]:
        """
        Returns all `Tensor` or `PhiTreeNode` attribute names of `self` whose values are variable.
        Variables denote values that can change from one function call to the next or for which gradients can be recorded.
        If this method is not implemented, all attributes returned by `__value_attrs__()` are considered variable.

        The returned properties are used by the following functions:

        - `jit_compile()`
        - `jit_compile_linear()`
        - `stop_gradient()`
        - `jacobian()`
        - `custom_gradient()`

        Returns:
            `tuple` of `str` attributes.
                Calling `getattr(self, attr)` must return a `Tensor` or `PhiTreeNode` for all returned attributes.
        """
        raise NotImplementedError

    def __with_attrs__(self, **attrs):
        """
        Used by `phiml.math.copy_with`.
        Create a copy of this object which has the `Tensor` or `PhiTreeNode` attributes contained in `attrs` replaced.
        If this method is not implemented, tensor attributes are replaced using `setattr()`.

        Args:
            **attrs: `dict` mapping `str` attribute names to `Tensor` or `PhiTreeNode`.

        Returns:
            Altered copy of `self`
        """
        raise NotImplementedError

    def __eq__(self, other):
        raise NotImplementedError



class BoundDim:
    """
    Represents a dimension of a sliceable object to make slicing, renaming and retyping prettier.
    Any instance of `BoundDim` is bound to the sliceable object and is immutable.
    All operations upon the dim affect return a copy of the sliceable object.

    `BoundDim` objects are generally created by and for objects that are `Sliceable` (and therefore also `Shaped`).
    These objects should declare the following method to support the `.dim` syntax:

    ```python
    from phiml.math.magic import BoundDim

    class MyClass:

        def __getattr__(self, name: str) -> BoundDim:
            return BoundDim(self, name)
    ```

    **Usage**

    * `obj.dim.size` returns the dimension size.
    * `obj.dim.item_names` returns the dimension item names.
    * `obj.dim.exists` checks whether a dimension is listed in the shape of the bound object.
    * `obj.dim[0]` picks the first element along `dim`. The shape of the result will not contain `dim`.
    * `obj.dim[1:-1]` discards the first and last element along `dim`.
    * `obj.dim.rename('new_name')` renames `dim` to `new_name`.
    * `obj.dim.as_channel()` changes the type of `dim` to *channel*.
    * `obj.dim.unstack()` un-stacks the bound value along `dim`.
    * `for slice in obj.dim` loops over all slices of `dim`.

    Multiple dimensions can also be chained together using `obj.dim1.dim2...`.
    This supports the following operations:

    * `obj.dim1.dim2...volume` returns the product of the sizes
    * `obj.dim1.dim2...[0, -1]` takes the first element along `dim1` and the last element along `dim2`
    * `obj.dim1.dim2...pack(new_dim)` packs the dimensions into a new dimension with size equal to their volume
    * `obj.dim1.dim2...unstack()` un-stacks `obj` along multiple dimensions
    * `obj.dim1.dim2...retype(type)` Changes the type of all selected dimensions
    * `for slice in obj.dim1.dim2...` loops over all slices as if unstacking first
    """

    def __init__(self, obj, name: str):
        """
        Args:
            obj: `Sliceable` bound object.
            name: Dimension name as `str`.
        """
        if name.startswith('_') or ',' in name or ' ' in name:
            raise AttributeError(f"'{type(self)}' object has no attribute '{name}'")
        if name == 'shape':
            raise AttributeError(f"Object of type {type(obj)} has no shape")
        if name == 'is_tensor_like':
            raise AttributeError(f"Trying to access {type(obj)}.is_tensor_like which does not exist. This is likely a TensorFlow JIT compile check. This name is reserved.")
        assert isinstance(obj, Sliceable) and isinstance(obj, Shaped), f"Cannot create BoundDim for {type(obj).__name__}. Objects must be Sliceable and Shaped, see https://tum-pbs.github.io/PhiML/phiml/math/magic.html"
        self.obj = obj
        self.name = name

    @property
    def dual(self):
        return BoundDim(self.obj, '~' + self.name)

    @property
    def exists(self):
        """ Whether the dimension is listed in the `Shape` of the object. """
        return self.name in shape(self.obj)

    def __repr__(self):
        if self.name not in shape(self.obj):
            return f"{type(self.obj).__name__}.{self.name} (non-existent)"
        items = self.item_names
        if items is not None:
            if len(items) <= 4:
                size_repr = ",".join(items)
            else:
                size_repr = f"{self.size}:{items[0]}..{items[-1]}"
        else:
            size_repr = self.size
        from ._shape import SUPERSCRIPT
        return f"{type(self.obj).__name__}.{self.name}{SUPERSCRIPT.get(self.type.__name__, '?')}={size_repr}"

    @property
    def size(self):
        """ Length of this dimension as listed in the `Shape` of the bound object. """
        return shape(self.obj).get_size(self.name) if self.exists else None

    volume = size

    @property
    def size_or_1(self):
        return shape(self.obj).get_size(self.name) if self.exists else 1

    @property
    def type(self) -> Callable:
        """
        The dimension type of this bound dimension. Must be one of `batch`, `spatial`, `instance`, `channel`.

        Returns:

        """
        return shape(self.obj).get_dim_type(self.name)

    @property
    def item_names(self):
        return shape(self.obj).get_item_names(self.name)

    def __getitem__(self, item):
        return self.obj[{self.name: item}]

    def __setitem__(self, key, value):
        self.obj[{self.name: key}] = value

    def __getattr__(self, item):
        return _BoundDims(self.obj, (self.name, item))

    def keys(self):
        """
        Allows unstacking with item names as `dict(**obj.dim)`.

        Returns:
            Sequence of item names or indices.
        """
        if not self.exists:
            raise SyntaxError(f"Cannot get keys of nonexistent dimension {self}")
        if self.item_names is not None:
            return self.item_names
        else:
            return range(self.size)

    def unstack(self, size: Union[int, None] = None) -> tuple:
        """
        Lists the slices along this dimension as a `tuple`.

        Args:
            size: (optional) If given as `int`, this dimension can be unstacked even if it is not present on the object.
                In that case, `size` copies of the object are returned.

        Returns:
            `tuple` of `Sliceable`
        """
        from ._magic_ops import unstack
        if size is None:
            return unstack(self.obj, self.name)
        else:
            if self.exists:
                unstacked = unstack(self.obj, self.name)
                assert len(unstacked) == size, f"Size of dimension {self.name} does not match {size}."
                return unstacked
            else:
                return (self.obj,) * size

    def __iter__(self):
        """ Iterate over slices along this dim """
        if self.exists:
            return iter(self.unstack())
        else:
            return iter([self.obj])

    def __call__(self, *args, **kwargs):
        raise TypeError(f"Method {type(self.obj).__name__}.{self.name}() does not exist.")

    def rename(self, name: str, **kwargs):
        """
        Returns a shallow copy of the `Tensor` where this dimension has the specified name.

        See Also:
            `phiml.math.rename_dims()`
        """
        if not self.exists:
            return self.obj
        from ._magic_ops import rename_dims
        return rename_dims(self.obj, self.name, name, **kwargs)

    def retype(self, dim_type: Callable, **kwargs):
        """
        Returns a shallow copy of the `Tensor` where this dimension has the specified type.

        See Also:
            `phiml.math.rename_dims()`
        """
        new_dim = dim_type(**{self.name: self.item_names or self.size})
        from ._magic_ops import rename_dims
        return rename_dims(self.obj, self.name, new_dim, **kwargs)

    def as_batch(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *batch*. """
        if not self.exists:
            return self.obj
        return self.retype(batch) if name is None else self.replace(batch(name=self.item_names or self.size))

    def as_spatial(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *spatial*. """
        if not self.exists:
            return self.obj
        return self.retype(spatial) if name is None else self.replace(spatial(name=self.item_names or self.size))

    def as_channel(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *channel*. """
        if not self.exists:
            return self.obj
        return self.retype(channel) if name is None else self.replace(channel(name=self.item_names or self.size))

    def as_instance(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *instance*. """
        if not self.exists:
            return self.obj
        return self.retype(instance) if name is None else self.replace(instance(name=self.item_names or self.size))

    def as_dual(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *instance*. """
        if not self.exists:
            return self.obj
        return self.retype(dual) if name is None else self.replace(dual(name=self.item_names or self.size))

    def replace(self, dim: Shape, **kwargs):
        """
        Returns a shallow copy of the `Tensor` where this dimension has been replaced by `dim`.

        See Also:
            `phiml.math.rename_dims()`
        """
        from ._magic_ops import rename_dims
        return rename_dims(self.obj, self.name, dim, **kwargs)

    def unpack(self, *dims: Shape, **kwargs):
        """
        Returns a shallow copy of the `Tensor` where this dimension has been unpacked into `dims`.

        See Also:
            `phiml.math.unpack_dim()`
        """
        from ._magic_ops import unpack_dim
        return unpack_dim(self.obj, self.name, *dims, **kwargs)

    def __bool__(self):
        raise SyntaxError(f"{self} cannot be converted to bool. The property you want to access likely does not exist.")


class _BoundDims:

    def __init__(self, obj, dims: Tuple[str, ...]):
        self.obj = obj
        self.dims = dims

    @property
    def dual(self):
        last_dual = "~" + self.dims[-1]
        return _BoundDims(self.obj, self.dims[:-1] + (last_dual,))

    def __getitem__(self, item):
        assert isinstance(item, tuple), f"A tuple of slices is required for slicing multiple dimensions at once but got {type(item)}"
        assert len(item) == len(self.dims), f"Number of slices must equal number of dimensions but got {len(item)} for dims {self.dims}"
        return self.obj[{dim: i for dim, i in zip(self.dims, item)}]

    def __getattr__(self, item):
        if len(self.dims) > 10:  # to avoid recursion limit
            raise RuntimeError("Maximum BoundDim length reached")
        return _BoundDims(self.obj, self.dims + (item,))

    def __len__(self):
        return self.volume

    @property
    def size(self):
        raise SyntaxError("dim.size only exists for single dimensions. Use .volume for multiple dimensions")

    @property
    def volume(self):
        return shape(self.obj).only(self.dims).volume

    def pack(self, packed_dim: Shape, pos=None, **kwargs):
        from ._magic_ops import pack_dims
        return pack_dims(self.obj, self.dims, packed_dim, pos=pos, **kwargs)

    def unstack(self) -> tuple:
        from ._magic_ops import unstack
        return unstack(self.obj, self.dims)

    def __iter__(self):
        """ Iterate over slices along this dim """
        return iter(self.unstack())

    def __call__(self, *args, **kwargs):
        raise TypeError(f"Method {type(self.obj).__name__}.{self.name}() does not exist.")

    def __repr__(self):
        parts = [type(self.obj).__name__, *self.dims]
        return ".".join(parts)

    def retype(self, dim_type: Callable, **kwargs):
        """
        Returns a shallow copy of the `Tensor` where this dimension has the specified type.

        See Also:
            `phiml.math.rename_dims()`
        """
        s = shape(self.obj)
        new_dims = concat_shapes(*[dim_type(**{dim: s.get_item_names(dim) or s.get_size(dim)}) for dim in self.dims])
        from ._magic_ops import rename_dims
        return rename_dims(self.obj, self.dims, new_dims, **kwargs)

    def as_batch(self):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *batch*. """
        return self.retype(batch)

    def as_spatial(self):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *spatial*. """
        return self.retype(spatial)

    def as_channel(self):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *channel*. """
        return self.retype(channel)

    def as_instance(self):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *instance*. """
        return self.retype(instance)

    def as_dual(self):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *instance*. """
        return self.retype(dual)

    def __bool__(self):
        raise SyntaxError(f"{self} cannot be converted to bool. The property you want to access likely does not exist.")


def slicing_dict(obj, item) -> dict:
    """
    Creates a slicing `dict` from `item` where `item` is an arbitrary value passed to `__getitem__()`.

    `Sliceable` objects should call this function inside `__getitem__()`, passing `self` and `item`.

    Args:
        obj: Object to be sliced.
        item: Slices.

    Returns:
        `dict` mapping dimension names to slices.
    """
    if isinstance(item, Shape):
        item = item.name_list
    if isinstance(item, dict):
        def valid_key(key):
            if callable(key):
                key = key(obj)
            if isinstance(key, str):
                return key
            elif isinstance(key, Shape):
                return key.name
            else:
                raise ValueError(f"Illegal slicing dim: {key}. Only str and single-dim Shape and filters are allowed. While trying to slice {obj} by {item}")
        return {valid_key(k): v for k, v in item.items()}
    elif isinstance(item, tuple):
        if item[0] == Ellipsis:
            assert len(item) - 1 == shape(obj).channel_rank
            return {name: selection for name, selection in zip(channel(obj).names, item[1:])}
        elif len(item) == shape(obj).channel_rank:
            if len(item) > 1:
                warnings.warn("NumPy-style slicing for more than one channel dimension is highly discouraged. Use a dict or the special slicing syntax value.dim[slice] instead. See https://tum-pbs.github.io/PhiML/Introduction.html#Slicing", SyntaxWarning, stacklevel=3)
            return {name: selection for name, selection in zip(channel(obj).names, item)}
        elif shape(obj).channel_rank == 1 and all(isinstance(e, str) for e in item):
            return {channel(obj).name: item}
        else:
            raise AssertionError(f"Cannot slice {obj}[{item}]. Use a dict or the special slicing syntax value.dim[slice] instead. See https://tum-pbs.github.io/PhiML/Introduction.html#Slicing")
    else:
        if shape(obj).channel_rank == 1:
            return {channel(obj).name: item}
        elif non_batch(obj).rank == 1:
            return {non_batch(obj).name: item}
        else:
            from ._tensors import Tensor
            class_name = "Tensor" if isinstance(obj, Tensor) else type(obj).__name__
            raise AssertionError(f"Slicing {class_name}[{type(item).__name__}] is only supported for 1D values (excluding batch dimensions) but shape is {shape(obj)}")


class OtherMagicFunctions:

    def __cast__(self, dtype: DType):
        raise NotImplementedError


__pdoc__ = {}  # Show all magic functions in pdoc3
for cls_name, cls in dict(globals()).items():
    if isinstance(cls, type) and type(cls) != type and not cls_name.startswith('_'):
        for magic_function in dir(cls):
            if magic_function.startswith('__') and magic_function.endswith('__') and not hasattr(object, magic_function) and magic_function != '__weakref__':
                __pdoc__[f'{cls_name}.{magic_function}'] = True

Functions

def slicing_dict(obj, item) ‑> dict

Creates a slicing dict from item where item is an arbitrary value passed to __getitem__().

Sliceable objects should call this function inside __getitem__(), passing self and item.

Args

obj

Object to be sliced.

item

Slices.

Returns

dict mapping dimension names to slices.

Expand source code

def slicing_dict(obj, item) -> dict:
    """
    Creates a slicing `dict` from `item` where `item` is an arbitrary value passed to `__getitem__()`.

    `Sliceable` objects should call this function inside `__getitem__()`, passing `self` and `item`.

    Args:
        obj: Object to be sliced.
        item: Slices.

    Returns:
        `dict` mapping dimension names to slices.
    """
    if isinstance(item, Shape):
        item = item.name_list
    if isinstance(item, dict):
        def valid_key(key):
            if callable(key):
                key = key(obj)
            if isinstance(key, str):
                return key
            elif isinstance(key, Shape):
                return key.name
            else:
                raise ValueError(f"Illegal slicing dim: {key}. Only str and single-dim Shape and filters are allowed. While trying to slice {obj} by {item}")
        return {valid_key(k): v for k, v in item.items()}
    elif isinstance(item, tuple):
        if item[0] == Ellipsis:
            assert len(item) - 1 == shape(obj).channel_rank
            return {name: selection for name, selection in zip(channel(obj).names, item[1:])}
        elif len(item) == shape(obj).channel_rank:
            if len(item) > 1:
                warnings.warn("NumPy-style slicing for more than one channel dimension is highly discouraged. Use a dict or the special slicing syntax value.dim[slice] instead. See https://tum-pbs.github.io/PhiML/Introduction.html#Slicing", SyntaxWarning, stacklevel=3)
            return {name: selection for name, selection in zip(channel(obj).names, item)}
        elif shape(obj).channel_rank == 1 and all(isinstance(e, str) for e in item):
            return {channel(obj).name: item}
        else:
            raise AssertionError(f"Cannot slice {obj}[{item}]. Use a dict or the special slicing syntax value.dim[slice] instead. See https://tum-pbs.github.io/PhiML/Introduction.html#Slicing")
    else:
        if shape(obj).channel_rank == 1:
            return {channel(obj).name: item}
        elif non_batch(obj).rank == 1:
            return {non_batch(obj).name: item}
        else:
            from ._tensors import Tensor
            class_name = "Tensor" if isinstance(obj, Tensor) else type(obj).__name__
            raise AssertionError(f"Slicing {class_name}[{type(item).__name__}] is only supported for 1D values (excluding batch dimensions) but shape is {shape(obj)}")

Classes

class BoundDim (obj, name: str)

Represents a dimension of a sliceable object to make slicing, renaming and retyping prettier. Any instance of BoundDim is bound to the sliceable object and is immutable. All operations upon the dim affect return a copy of the sliceable object.

BoundDim objects are generally created by and for objects that are Sliceable (and therefore also Shaped). These objects should declare the following method to support the .dim syntax:

from phiml.math.magic import BoundDim

class MyClass:

    def __getattr__(self, name: str) -> BoundDim:
        return BoundDim(self, name)

Usage

• obj.dim.size returns the dimension size.
• obj.dim.item_names returns the dimension item names.
• obj.dim.exists checks whether a dimension is listed in the shape of the bound object.
• obj.dim[0] picks the first element along dim. The shape of the result will not contain dim.
• obj.dim[1:-1] discards the first and last element along dim.
• obj.dim.rename('new_name') renames dim to new_name.
• obj.dim.as_channel() changes the type of dim to channel.
• obj.dim.unstack() un-stacks the bound value along dim.
• for slice in obj.dim loops over all slices of dim.

Multiple dimensions can also be chained together using obj.dim1.dim2…. This supports the following operations:

• obj.dim1.dim2…volume returns the product of the sizes
• obj.dim1.dim2...[0, -1] takes the first element along dim1 and the last element along dim2
• obj.dim1.dim2…pack(new_dim) packs the dimensions into a new dimension with size equal to their volume
• obj.dim1.dim2…unstack() un-stacks obj along multiple dimensions
• obj.dim1.dim2…retype(type) Changes the type of all selected dimensions
• for slice in obj.dim1.dim2… loops over all slices as if unstacking first

Args

obj

Sliceable bound object.

name

Dimension name as str.

Expand source code

class BoundDim:
    """
    Represents a dimension of a sliceable object to make slicing, renaming and retyping prettier.
    Any instance of `BoundDim` is bound to the sliceable object and is immutable.
    All operations upon the dim affect return a copy of the sliceable object.

    `BoundDim` objects are generally created by and for objects that are `Sliceable` (and therefore also `Shaped`).
    These objects should declare the following method to support the `.dim` syntax:

    ```python
    from phiml.math.magic import BoundDim

    class MyClass:

        def __getattr__(self, name: str) -> BoundDim:
            return BoundDim(self, name)
    ```

    **Usage**

    * `obj.dim.size` returns the dimension size.
    * `obj.dim.item_names` returns the dimension item names.
    * `obj.dim.exists` checks whether a dimension is listed in the shape of the bound object.
    * `obj.dim[0]` picks the first element along `dim`. The shape of the result will not contain `dim`.
    * `obj.dim[1:-1]` discards the first and last element along `dim`.
    * `obj.dim.rename('new_name')` renames `dim` to `new_name`.
    * `obj.dim.as_channel()` changes the type of `dim` to *channel*.
    * `obj.dim.unstack()` un-stacks the bound value along `dim`.
    * `for slice in obj.dim` loops over all slices of `dim`.

    Multiple dimensions can also be chained together using `obj.dim1.dim2...`.
    This supports the following operations:

    * `obj.dim1.dim2...volume` returns the product of the sizes
    * `obj.dim1.dim2...[0, -1]` takes the first element along `dim1` and the last element along `dim2`
    * `obj.dim1.dim2...pack(new_dim)` packs the dimensions into a new dimension with size equal to their volume
    * `obj.dim1.dim2...unstack()` un-stacks `obj` along multiple dimensions
    * `obj.dim1.dim2...retype(type)` Changes the type of all selected dimensions
    * `for slice in obj.dim1.dim2...` loops over all slices as if unstacking first
    """

    def __init__(self, obj, name: str):
        """
        Args:
            obj: `Sliceable` bound object.
            name: Dimension name as `str`.
        """
        if name.startswith('_') or ',' in name or ' ' in name:
            raise AttributeError(f"'{type(self)}' object has no attribute '{name}'")
        if name == 'shape':
            raise AttributeError(f"Object of type {type(obj)} has no shape")
        if name == 'is_tensor_like':
            raise AttributeError(f"Trying to access {type(obj)}.is_tensor_like which does not exist. This is likely a TensorFlow JIT compile check. This name is reserved.")
        assert isinstance(obj, Sliceable) and isinstance(obj, Shaped), f"Cannot create BoundDim for {type(obj).__name__}. Objects must be Sliceable and Shaped, see https://tum-pbs.github.io/PhiML/phiml/math/magic.html"
        self.obj = obj
        self.name = name

    @property
    def dual(self):
        return BoundDim(self.obj, '~' + self.name)

    @property
    def exists(self):
        """ Whether the dimension is listed in the `Shape` of the object. """
        return self.name in shape(self.obj)

    def __repr__(self):
        if self.name not in shape(self.obj):
            return f"{type(self.obj).__name__}.{self.name} (non-existent)"
        items = self.item_names
        if items is not None:
            if len(items) <= 4:
                size_repr = ",".join(items)
            else:
                size_repr = f"{self.size}:{items[0]}..{items[-1]}"
        else:
            size_repr = self.size
        from ._shape import SUPERSCRIPT
        return f"{type(self.obj).__name__}.{self.name}{SUPERSCRIPT.get(self.type.__name__, '?')}={size_repr}"

    @property
    def size(self):
        """ Length of this dimension as listed in the `Shape` of the bound object. """
        return shape(self.obj).get_size(self.name) if self.exists else None

    volume = size

    @property
    def size_or_1(self):
        return shape(self.obj).get_size(self.name) if self.exists else 1

    @property
    def type(self) -> Callable:
        """
        The dimension type of this bound dimension. Must be one of `batch`, `spatial`, `instance`, `channel`.

        Returns:

        """
        return shape(self.obj).get_dim_type(self.name)

    @property
    def item_names(self):
        return shape(self.obj).get_item_names(self.name)

    def __getitem__(self, item):
        return self.obj[{self.name: item}]

    def __setitem__(self, key, value):
        self.obj[{self.name: key}] = value

    def __getattr__(self, item):
        return _BoundDims(self.obj, (self.name, item))

    def keys(self):
        """
        Allows unstacking with item names as `dict(**obj.dim)`.

        Returns:
            Sequence of item names or indices.
        """
        if not self.exists:
            raise SyntaxError(f"Cannot get keys of nonexistent dimension {self}")
        if self.item_names is not None:
            return self.item_names
        else:
            return range(self.size)

    def unstack(self, size: Union[int, None] = None) -> tuple:
        """
        Lists the slices along this dimension as a `tuple`.

        Args:
            size: (optional) If given as `int`, this dimension can be unstacked even if it is not present on the object.
                In that case, `size` copies of the object are returned.

        Returns:
            `tuple` of `Sliceable`
        """
        from ._magic_ops import unstack
        if size is None:
            return unstack(self.obj, self.name)
        else:
            if self.exists:
                unstacked = unstack(self.obj, self.name)
                assert len(unstacked) == size, f"Size of dimension {self.name} does not match {size}."
                return unstacked
            else:
                return (self.obj,) * size

    def __iter__(self):
        """ Iterate over slices along this dim """
        if self.exists:
            return iter(self.unstack())
        else:
            return iter([self.obj])

    def __call__(self, *args, **kwargs):
        raise TypeError(f"Method {type(self.obj).__name__}.{self.name}() does not exist.")

    def rename(self, name: str, **kwargs):
        """
        Returns a shallow copy of the `Tensor` where this dimension has the specified name.

        See Also:
            `phiml.math.rename_dims()`
        """
        if not self.exists:
            return self.obj
        from ._magic_ops import rename_dims
        return rename_dims(self.obj, self.name, name, **kwargs)

    def retype(self, dim_type: Callable, **kwargs):
        """
        Returns a shallow copy of the `Tensor` where this dimension has the specified type.

        See Also:
            `phiml.math.rename_dims()`
        """
        new_dim = dim_type(**{self.name: self.item_names or self.size})
        from ._magic_ops import rename_dims
        return rename_dims(self.obj, self.name, new_dim, **kwargs)

    def as_batch(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *batch*. """
        if not self.exists:
            return self.obj
        return self.retype(batch) if name is None else self.replace(batch(name=self.item_names or self.size))

    def as_spatial(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *spatial*. """
        if not self.exists:
            return self.obj
        return self.retype(spatial) if name is None else self.replace(spatial(name=self.item_names or self.size))

    def as_channel(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *channel*. """
        if not self.exists:
            return self.obj
        return self.retype(channel) if name is None else self.replace(channel(name=self.item_names or self.size))

    def as_instance(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *instance*. """
        if not self.exists:
            return self.obj
        return self.retype(instance) if name is None else self.replace(instance(name=self.item_names or self.size))

    def as_dual(self, name: str = None):
        """ Returns a shallow copy of the `Tensor` where the type of this dimension is *instance*. """
        if not self.exists:
            return self.obj
        return self.retype(dual) if name is None else self.replace(dual(name=self.item_names or self.size))

    def replace(self, dim: Shape, **kwargs):
        """
        Returns a shallow copy of the `Tensor` where this dimension has been replaced by `dim`.

        See Also:
            `phiml.math.rename_dims()`
        """
        from ._magic_ops import rename_dims
        return rename_dims(self.obj, self.name, dim, **kwargs)

    def unpack(self, *dims: Shape, **kwargs):
        """
        Returns a shallow copy of the `Tensor` where this dimension has been unpacked into `dims`.

        See Also:
            `phiml.math.unpack_dim()`
        """
        from ._magic_ops import unpack_dim
        return unpack_dim(self.obj, self.name, *dims, **kwargs)

    def __bool__(self):
        raise SyntaxError(f"{self} cannot be converted to bool. The property you want to access likely does not exist.")

Subclasses

• phiml.math._tensors.TensorDim

Instance variables

var dual

Expand source code

@property
def dual(self):
    return BoundDim(self.obj, '~' + self.name)

var exists

Whether the dimension is listed in the Shape of the object.

Expand source code

@property
def exists(self):
    """ Whether the dimension is listed in the `Shape` of the object. """
    return self.name in shape(self.obj)

var item_names

Expand source code

@property
def item_names(self):
    return shape(self.obj).get_item_names(self.name)

var size

Length of this dimension as listed in the Shape of the bound object.

Expand source code

@property
def size(self):
    """ Length of this dimension as listed in the `Shape` of the bound object. """
    return shape(self.obj).get_size(self.name) if self.exists else None

var size_or_1

Expand source code

@property
def size_or_1(self):
    return shape(self.obj).get_size(self.name) if self.exists else 1

var type : Callable

The dimension type of this bound dimension. Must be one of batch, spatial, instance, channel.

Returns:

Expand source code

@property
def type(self) -> Callable:
    """
    The dimension type of this bound dimension. Must be one of `batch`, `spatial`, `instance`, `channel`.

    Returns:

    """
    return shape(self.obj).get_dim_type(self.name)

var volume

Length of this dimension as listed in the Shape of the bound object.

Expand source code

@property
def size(self):
    """ Length of this dimension as listed in the `Shape` of the bound object. """
    return shape(self.obj).get_size(self.name) if self.exists else None

Methods

def as_batch(self, name: str = None)

Returns a shallow copy of the Tensor where the type of this dimension is batch.

Expand source code

def as_batch(self, name: str = None):
    """ Returns a shallow copy of the `Tensor` where the type of this dimension is *batch*. """
    if not self.exists:
        return self.obj
    return self.retype(batch) if name is None else self.replace(batch(name=self.item_names or self.size))

def as_channel(self, name: str = None)

Returns a shallow copy of the Tensor where the type of this dimension is channel.

Expand source code

def as_channel(self, name: str = None):
    """ Returns a shallow copy of the `Tensor` where the type of this dimension is *channel*. """
    if not self.exists:
        return self.obj
    return self.retype(channel) if name is None else self.replace(channel(name=self.item_names or self.size))

def as_dual(self, name: str = None)

Returns a shallow copy of the Tensor where the type of this dimension is instance.

Expand source code

def as_dual(self, name: str = None):
    """ Returns a shallow copy of the `Tensor` where the type of this dimension is *instance*. """
    if not self.exists:
        return self.obj
    return self.retype(dual) if name is None else self.replace(dual(name=self.item_names or self.size))

def as_instance(self, name: str = None)

Returns a shallow copy of the Tensor where the type of this dimension is instance.

Expand source code

def as_instance(self, name: str = None):
    """ Returns a shallow copy of the `Tensor` where the type of this dimension is *instance*. """
    if not self.exists:
        return self.obj
    return self.retype(instance) if name is None else self.replace(instance(name=self.item_names or self.size))

def as_spatial(self, name: str = None)

Returns a shallow copy of the Tensor where the type of this dimension is spatial.

Expand source code

def as_spatial(self, name: str = None):
    """ Returns a shallow copy of the `Tensor` where the type of this dimension is *spatial*. """
    if not self.exists:
        return self.obj
    return self.retype(spatial) if name is None else self.replace(spatial(name=self.item_names or self.size))

def keys(self)

Allows unstacking with item names as dict(**obj.dim).

Returns

Sequence of item names or indices.

Expand source code

def keys(self):
    """
    Allows unstacking with item names as `dict(**obj.dim)`.

    Returns:
        Sequence of item names or indices.
    """
    if not self.exists:
        raise SyntaxError(f"Cannot get keys of nonexistent dimension {self}")
    if self.item_names is not None:
        return self.item_names
    else:
        return range(self.size)

def rename(self, name: str, **kwargs)

Returns a shallow copy of the Tensor where this dimension has the specified name.

See Also: rename_dims()

Expand source code

def rename(self, name: str, **kwargs):
    """
    Returns a shallow copy of the `Tensor` where this dimension has the specified name.

    See Also:
        `phiml.math.rename_dims()`
    """
    if not self.exists:
        return self.obj
    from ._magic_ops import rename_dims
    return rename_dims(self.obj, self.name, name, **kwargs)

def replace(self, dim: phiml.math._shape.Shape, **kwargs)

Returns a shallow copy of the Tensor where this dimension has been replaced by dim.

See Also: rename_dims()

Expand source code

def replace(self, dim: Shape, **kwargs):
    """
    Returns a shallow copy of the `Tensor` where this dimension has been replaced by `dim`.

    See Also:
        `phiml.math.rename_dims()`
    """
    from ._magic_ops import rename_dims
    return rename_dims(self.obj, self.name, dim, **kwargs)

def retype(self, dim_type: Callable, **kwargs)

Returns a shallow copy of the Tensor where this dimension has the specified type.

See Also: rename_dims()

Expand source code

def retype(self, dim_type: Callable, **kwargs):
    """
    Returns a shallow copy of the `Tensor` where this dimension has the specified type.

    See Also:
        `phiml.math.rename_dims()`
    """
    new_dim = dim_type(**{self.name: self.item_names or self.size})
    from ._magic_ops import rename_dims
    return rename_dims(self.obj, self.name, new_dim, **kwargs)

def unpack(self, *dims: phiml.math._shape.Shape, **kwargs)

Returns a shallow copy of the Tensor where this dimension has been unpacked into dims.

See Also: unpack_dim()

Expand source code

def unpack(self, *dims: Shape, **kwargs):
    """
    Returns a shallow copy of the `Tensor` where this dimension has been unpacked into `dims`.

    See Also:
        `phiml.math.unpack_dim()`
    """
    from ._magic_ops import unpack_dim
    return unpack_dim(self.obj, self.name, *dims, **kwargs)

def unstack(self, size: Optional[int] = None) ‑> tuple

Lists the slices along this dimension as a tuple.

Args

size

(optional) If given as int, this dimension can be unstacked even if it is not present on the object. In that case, size copies of the object are returned.

Returns

tuple of Sliceable

Expand source code

def unstack(self, size: Union[int, None] = None) -> tuple:
    """
    Lists the slices along this dimension as a `tuple`.

    Args:
        size: (optional) If given as `int`, this dimension can be unstacked even if it is not present on the object.
            In that case, `size` copies of the object are returned.

    Returns:
        `tuple` of `Sliceable`
    """
    from ._magic_ops import unstack
    if size is None:
        return unstack(self.obj, self.name)
    else:
        if self.exists:
            unstacked = unstack(self.obj, self.name)
            assert len(unstacked) == size, f"Size of dimension {self.name} does not match {size}."
            return unstacked
        else:
            return (self.obj,) * size

class OtherMagicFunctions

Expand source code

class OtherMagicFunctions:

    def __cast__(self, dtype: DType):
        raise NotImplementedError

class PhiTreeNode

Φ-tree nodes can be iterated over and disassembled or flattened into elementary objects, such as tensors. Tensor instances as well as PyTree nodes (tuple, list, dict with str keys) are Φ-tree nodes. All data classes are also considered PhiTreeNodes as of version 2.3.

For custom classes to be considered Φ-tree nodes, they have to be a dataclass or implement one of the following magic methods:

• __variable_attrs__()
• __value_attrs__()

Dataclasses may also implement these functions to specify which attributes should be considered value / variable properties.

Additionally, Φ-tree nodes must override __eq__() to allow comparison of data-stripped (key) instances.

To check whether an object is a Φ-tree node, use isinstance(obj, PhiTreeNode).

Usage in phiml.math:

Φ-tree nodes can be used as keys, for example in jit_compile(). They are converted to keys by stripping all variable tensors and replacing them by a placeholder object. In key mode, __eq__() compares all non-variable properties that might invalidate a trace when changed.

Disassembly and assembly of Φ-tree nodes uses replace() which will call __with_attrs__ if implemented.

Expand source code

class PhiTreeNode(metaclass=_PhiTreeNodeType):
    """
    Φ-tree nodes can be iterated over and disassembled or flattened into elementary objects, such as tensors.
    `phiml.math.Tensor` instances as well as PyTree nodes (`tuple`, `list`, `dict` with `str` keys) are Φ-tree nodes.
    All data classes are also considered PhiTreeNodes as of version 2.3.

    For custom classes to be considered Φ-tree nodes, they have to be a dataclass or implement one of the following magic methods:

    * `__variable_attrs__()`
    * `__value_attrs__()`

    Dataclasses may also implement these functions to specify which attributes should be considered value / variable properties.

    Additionally, Φ-tree nodes must override `__eq__()` to allow comparison of data-stripped (key) instances.

    To check whether an object is a Φ-tree node, use `isinstance(obj, PhiTreeNode)`.

    **Usage in `phiml.math`:**

    Φ-tree nodes can be used as keys, for example in `jit_compile()`.
    They are converted to keys by stripping all variable tensors and replacing them by a placeholder object.
    In key mode, `__eq__()` compares all non-variable properties that might invalidate a trace when changed.

    Disassembly and assembly of Φ-tree nodes uses `phiml.math.copy_with` which will call `__with_attrs__` if implemented.
    """

    def __value_attrs__(self) -> Tuple[str, ...]:
        """
        Returns all `Tensor` or `PhiTreeNode` attribute names of `self` that should be transformed by single-operand math operations,
        such as `sin()`, `exp()`.

        Returns:
            `tuple` of `str` attributes.
                Calling `getattr(self, attr)` must return a `Tensor` or `PhiTreeNode` for all returned attributes.
        """
        raise NotImplementedError

    def __variable_attrs__(self) -> Tuple[str, ...]:
        """
        Returns all `Tensor` or `PhiTreeNode` attribute names of `self` whose values are variable.
        Variables denote values that can change from one function call to the next or for which gradients can be recorded.
        If this method is not implemented, all attributes returned by `__value_attrs__()` are considered variable.

        The returned properties are used by the following functions:

        - `jit_compile()`
        - `jit_compile_linear()`
        - `stop_gradient()`
        - `jacobian()`
        - `custom_gradient()`

        Returns:
            `tuple` of `str` attributes.
                Calling `getattr(self, attr)` must return a `Tensor` or `PhiTreeNode` for all returned attributes.
        """
        raise NotImplementedError

    def __with_attrs__(self, **attrs):
        """
        Used by `phiml.math.copy_with`.
        Create a copy of this object which has the `Tensor` or `PhiTreeNode` attributes contained in `attrs` replaced.
        If this method is not implemented, tensor attributes are replaced using `setattr()`.

        Args:
            **attrs: `dict` mapping `str` attribute names to `Tensor` or `PhiTreeNode`.

        Returns:
            Altered copy of `self`
        """
        raise NotImplementedError

    def __eq__(self, other):
        raise NotImplementedError

Methods

def __value_attrs__(self) ‑> Tuple[str, ...]

Returns all Tensor or PhiTreeNode attribute names of self that should be transformed by single-operand math operations, such as sin(), exp().

Returns

tuple of str attributes. Calling getattr(self, attr) must return a Tensor or PhiTreeNode for all returned attributes.

Expand source code

def __value_attrs__(self) -> Tuple[str, ...]:
    """
    Returns all `Tensor` or `PhiTreeNode` attribute names of `self` that should be transformed by single-operand math operations,
    such as `sin()`, `exp()`.

    Returns:
        `tuple` of `str` attributes.
            Calling `getattr(self, attr)` must return a `Tensor` or `PhiTreeNode` for all returned attributes.
    """
    raise NotImplementedError

def __variable_attrs__(self) ‑> Tuple[str, ...]

Returns all Tensor or PhiTreeNode attribute names of self whose values are variable. Variables denote values that can change from one function call to the next or for which gradients can be recorded. If this method is not implemented, all attributes returned by __value_attrs__() are considered variable.

The returned properties are used by the following functions:

• jit_compile()
• jit_compile_linear()
• stop_gradient()
• jacobian()
• custom_gradient()

Returns

tuple of str attributes. Calling getattr(self, attr) must return a Tensor or PhiTreeNode for all returned attributes.

Expand source code

def __variable_attrs__(self) -> Tuple[str, ...]:
    """
    Returns all `Tensor` or `PhiTreeNode` attribute names of `self` whose values are variable.
    Variables denote values that can change from one function call to the next or for which gradients can be recorded.
    If this method is not implemented, all attributes returned by `__value_attrs__()` are considered variable.

    The returned properties are used by the following functions:

    - `jit_compile()`
    - `jit_compile_linear()`
    - `stop_gradient()`
    - `jacobian()`
    - `custom_gradient()`

    Returns:
        `tuple` of `str` attributes.
            Calling `getattr(self, attr)` must return a `Tensor` or `PhiTreeNode` for all returned attributes.
    """
    raise NotImplementedError

def __with_attrs__(self, **attrs)

Used by replace(). Create a copy of this object which has the Tensor or PhiTreeNode attributes contained in attrs replaced. If this method is not implemented, tensor attributes are replaced using setattr().

Args

**attrs

dict mapping str attribute names to Tensor or PhiTreeNode.

Returns

Altered copy of self

Expand source code

def __with_attrs__(self, **attrs):
    """
    Used by `phiml.math.copy_with`.
    Create a copy of this object which has the `Tensor` or `PhiTreeNode` attributes contained in `attrs` replaced.
    If this method is not implemented, tensor attributes are replaced using `setattr()`.

    Args:
        **attrs: `dict` mapping `str` attribute names to `Tensor` or `PhiTreeNode`.

    Returns:
        Altered copy of `self`
    """
    raise NotImplementedError

class Shapable

Shapable objects can be stacked, concatenated and reshaped.

To be considered Shapable, objects must be Sliceable and Shaped and implement

• __stack__() or
• __concat__() and __expand__().

Objects should additionally implement the other magic methods for performance reasons.

Usage in phiml.math:

Shapable objects can be used with the following functions in addition to what they inherit from being Sliceable and Shaped:

• stack()
• concat()
• expand()
• rename_dims()
• pack_dims()
• unpack_dim()
• flatten()

Additionally, the phiml.math.BoundDim syntax for dimension renaming and retyping is enabled, e.g. obj.dim.as_channel('vector').

Expand source code

class Shapable(metaclass=_ShapableType):
    """
    Shapable objects can be stacked, concatenated and reshaped.

    To be considered `Shapable`, objects must be `Sliceable` and `Shaped` and implement

    * `__stack__()` or
    * `__concat__()` and `__expand__()`.

    Objects should additionally implement the other magic methods for performance reasons.

    **Usage in `phiml.math`:**

    Shapable objects can be used with the following functions in addition to what they inherit from being `Sliceable` and `Shaped`:

    * `phiml.math.stack`
    * `phiml.math.concat`
    * `phiml.math.expand`
    * `phiml.math.rename_dims`
    * `phiml.math.pack_dims`
    * `phiml.math.unpack_dim`
    * `phiml.math.flatten`

    Additionally, the `phiml.math.BoundDim` syntax for dimension renaming and retyping is enabled, e.g. `obj.dim.as_channel('vector')`.
    """
    @staticmethod
    def __stack__(values: tuple, dim: Shape, **kwargs) -> 'Shapable':
        """
        Stack all `values` into a single instance along the new dimension `dim`.

        This method can be implemented as a bound method or as a `staticmethod` (without the `self` argument).

        Args:
            values: `tuple` of `Shapable` objects to be stacked. `self` is included in that list at least once.
            dim: Single-dimension `Shape`. This dimension must not be present with any of the `values`.
                The dimension fulfills the condition `dim.size == len(values)`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` representing the stacked slices.
            Its shape includes `dim` in addition to the dimensions present in `values`.
            If such a representation cannot be created because some values in `values` are not supported, returns `NotImplemented`.
        """
        raise NotImplementedError

    @staticmethod
    def __concat__(values: tuple, dim: str, **kwargs) -> 'Shapable':
        """
        Concatenate `values` along `dim`.

        This method can be implemented as a bound method or as a `staticmethod` (without the `self` argument).

        Args:
            values: Values to concatenate. `self` is included in that list at least once.
            dim: Dimension nams as `str`, must be present in all `values`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` representing the concatenated values or `NotImplemented` to revert to default behavior for this object.
            When returning a valid object, the size of `dim` must be equal to the sum of all `dim` sizes in `values`.
            If such a representation cannot be created because some values in `values` are not supported, returns `NotImplemented`.
        """
        raise NotImplementedError

    def __expand__(self, dims: Shape, **kwargs) -> 'Shapable':
        """
        Adds new dimensions to this object.
        The value of this object is constant along the new dimensions.

        Args:
            dims: Dimensions to add.
                They are guaranteed to not already be present in `shape(self)`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

    def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Shapable':
        """
        Exchange existing dimensions.
        This can be used to rename dimensions, change dimension types or change item names.

        Args:
            dims: Dimensions to be replaced.
            new_dims: Replacement dimensions as `Shape` with `rank == len(dims)`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

    def __pack_dims__(self, dims: Tuple[str, ...], packed_dim: Shape, pos: Union[int, None], **kwargs) -> 'Shapable':
        """
        Compresses multiple dimensions into a single dimension by concatenating the elements.
        Elements along the new dimensions are laid out according to the order of `dims`.

        The type of the new dimension will be equal to the types of `dims`.
        If `dims` have varying types, the new dimension will be a batch dimension.

        Args:
            dims: Dimensions to be compressed in the specified order.
            packed_dim: Single-dimension `Shape`.
            pos: Index of new dimension. `None` for automatic, `-1` for last, `0` for first.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

    def __unpack_dim__(self, dim: str, unpacked_dims: Shape, **kwargs) -> 'Shapable':
        """
        Decompresses a tensor dimension by unstacking the elements along it.
        The compressed dimension `dim` is assumed to contain elements laid out according to the order of `unpacked_dims`.

        Args:
            dim: Dimension to be decompressed.
            unpacked_dims: `Shape`: Ordered dimensions to replace `dim`, fulfilling `unpacked_dims.volume == shape(self)[dim].rank`.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

    def __flatten__(self, flat_dim: Shape, flatten_batch: bool, **kwargs) -> 'Shapable':
        """
        Lays out all elements along a single dimension.
        This is equivalent to packing all dimensions.

        Args:
            flat_dim: Single dimension as `Shape`.
            flatten_batch: Whether to flatten batch dimensions as well.
            If `False`, batch dimensions are kept, only onn-batch dimensions are flattened.
            **kwargs: Additional keyword arguments required by specific implementations.
                Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
                Adding batch dimensions must always work without keyword arguments.

        Returns:
            New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

Static methods

def __concat__(values: tuple, dim: str, **kwargs) ‑> Shapable

Concatenate values along dim.

This method can be implemented as a bound method or as a staticmethod (without the self argument).

Args

values

Values to concatenate. self is included in that list at least once.

dim

Dimension nams as str, must be present in all values.

**kwargs

Additional keyword arguments required by specific implementations. Adding spatial dimensions to fields requires the bounds: Box argument specifying the physical extent of the new dimensions. Adding batch dimensions must always work without keyword arguments.

Returns

New instance of Shapable representing the concatenated values or NotImplemented to revert to default behavior for this object. When returning a valid object, the size of dim must be equal to the sum of all dim sizes in values. If such a representation cannot be created because some values in values are not supported, returns NotImplemented.

Expand source code

@staticmethod
def __concat__(values: tuple, dim: str, **kwargs) -> 'Shapable':
    """
    Concatenate `values` along `dim`.

    This method can be implemented as a bound method or as a `staticmethod` (without the `self` argument).

    Args:
        values: Values to concatenate. `self` is included in that list at least once.
        dim: Dimension nams as `str`, must be present in all `values`.
        **kwargs: Additional keyword arguments required by specific implementations.
            Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
            Adding batch dimensions must always work without keyword arguments.

    Returns:
        New instance of `Shapable` representing the concatenated values or `NotImplemented` to revert to default behavior for this object.
        When returning a valid object, the size of `dim` must be equal to the sum of all `dim` sizes in `values`.
        If such a representation cannot be created because some values in `values` are not supported, returns `NotImplemented`.
    """
    raise NotImplementedError

def __stack__(values: tuple, dim: phiml.math._shape.Shape, **kwargs) ‑> Shapable

Stack all values into a single instance along the new dimension dim.

This method can be implemented as a bound method or as a staticmethod (without the self argument).

Args

values

tuple of Shapable objects to be stacked. self is included in that list at least once.

dim

Single-dimension Shape. This dimension must not be present with any of the values. The dimension fulfills the condition dim.size == len(values).

**kwargs

Additional keyword arguments required by specific implementations. Adding spatial dimensions to fields requires the bounds: Box argument specifying the physical extent of the new dimensions. Adding batch dimensions must always work without keyword arguments.

Returns

New instance of Shapable representing the stacked slices. Its shape includes dim in addition to the dimensions present in values. If such a representation cannot be created because some values in values are not supported, returns NotImplemented.

Expand source code

@staticmethod
def __stack__(values: tuple, dim: Shape, **kwargs) -> 'Shapable':
    """
    Stack all `values` into a single instance along the new dimension `dim`.

    This method can be implemented as a bound method or as a `staticmethod` (without the `self` argument).

    Args:
        values: `tuple` of `Shapable` objects to be stacked. `self` is included in that list at least once.
        dim: Single-dimension `Shape`. This dimension must not be present with any of the `values`.
            The dimension fulfills the condition `dim.size == len(values)`.
        **kwargs: Additional keyword arguments required by specific implementations.
            Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
            Adding batch dimensions must always work without keyword arguments.

    Returns:
        New instance of `Shapable` representing the stacked slices.
        Its shape includes `dim` in addition to the dimensions present in `values`.
        If such a representation cannot be created because some values in `values` are not supported, returns `NotImplemented`.
    """
    raise NotImplementedError

Methods

def __expand__(self, dims: phiml.math._shape.Shape, **kwargs) ‑> Shapable

Adds new dimensions to this object. The value of this object is constant along the new dimensions.

Args

dims

Dimensions to add. They are guaranteed to not already be present in shape(self).

**kwargs

Additional keyword arguments required by specific implementations. Adding spatial dimensions to fields requires the bounds: Box argument specifying the physical extent of the new dimensions. Adding batch dimensions must always work without keyword arguments.

Returns

New instance of Shapable or NotImplemented to revert to default behavior for this object.

Expand source code

def __expand__(self, dims: Shape, **kwargs) -> 'Shapable':
    """
    Adds new dimensions to this object.
    The value of this object is constant along the new dimensions.

    Args:
        dims: Dimensions to add.
            They are guaranteed to not already be present in `shape(self)`.
        **kwargs: Additional keyword arguments required by specific implementations.
            Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
            Adding batch dimensions must always work without keyword arguments.

    Returns:
        New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
    """
    raise NotImplementedError

def __flatten__(self, flat_dim: phiml.math._shape.Shape, flatten_batch: bool, **kwargs) ‑> Shapable

Lays out all elements along a single dimension. This is equivalent to packing all dimensions.

Args

flat_dim

Single dimension as Shape.

flatten_batch

Whether to flatten batch dimensions as well.

If False, batch dimensions are kept, only onn-batch dimensions are flattened.

**kwargs

Additional keyword arguments required by specific implementations. Adding spatial dimensions to fields requires the bounds: Box argument specifying the physical extent of the new dimensions. Adding batch dimensions must always work without keyword arguments.

Returns

New instance of Shapable or NotImplemented to revert to default behavior for this object.

Expand source code

def __flatten__(self, flat_dim: Shape, flatten_batch: bool, **kwargs) -> 'Shapable':
    """
    Lays out all elements along a single dimension.
    This is equivalent to packing all dimensions.

    Args:
        flat_dim: Single dimension as `Shape`.
        flatten_batch: Whether to flatten batch dimensions as well.
        If `False`, batch dimensions are kept, only onn-batch dimensions are flattened.
        **kwargs: Additional keyword arguments required by specific implementations.
            Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
            Adding batch dimensions must always work without keyword arguments.

    Returns:
        New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
    """
    raise NotImplementedError

def __pack_dims__(self, dims: Tuple[str, ...], packed_dim: phiml.math._shape.Shape, pos: Optional[int], **kwargs) ‑> Shapable

Compresses multiple dimensions into a single dimension by concatenating the elements. Elements along the new dimensions are laid out according to the order of dims.

The type of the new dimension will be equal to the types of dims. If dims have varying types, the new dimension will be a batch dimension.

Args

dims

Dimensions to be compressed in the specified order.

packed_dim

Single-dimension Shape.

pos

Index of new dimension. None for automatic, -1 for last, 0 for first.

**kwargs

Additional keyword arguments required by specific implementations. Adding spatial dimensions to fields requires the bounds: Box argument specifying the physical extent of the new dimensions. Adding batch dimensions must always work without keyword arguments.

Returns

New instance of Shapable or NotImplemented to revert to default behavior for this object.

Expand source code

def __pack_dims__(self, dims: Tuple[str, ...], packed_dim: Shape, pos: Union[int, None], **kwargs) -> 'Shapable':
    """
    Compresses multiple dimensions into a single dimension by concatenating the elements.
    Elements along the new dimensions are laid out according to the order of `dims`.

    The type of the new dimension will be equal to the types of `dims`.
    If `dims` have varying types, the new dimension will be a batch dimension.

    Args:
        dims: Dimensions to be compressed in the specified order.
        packed_dim: Single-dimension `Shape`.
        pos: Index of new dimension. `None` for automatic, `-1` for last, `0` for first.
        **kwargs: Additional keyword arguments required by specific implementations.
            Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
            Adding batch dimensions must always work without keyword arguments.

    Returns:
        New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
    """
    raise NotImplementedError

def __replace_dims__(self, dims: Tuple[str, ...], new_dims: phiml.math._shape.Shape, **kwargs) ‑> Shapable

Exchange existing dimensions. This can be used to rename dimensions, change dimension types or change item names.

Args

dims

Dimensions to be replaced.

new_dims

Replacement dimensions as Shape with rank == len(dims).

**kwargs

Additional keyword arguments required by specific implementations. Adding spatial dimensions to fields requires the bounds: Box argument specifying the physical extent of the new dimensions. Adding batch dimensions must always work without keyword arguments.

Returns

New instance of Shapable or NotImplemented to revert to default behavior for this object.

Expand source code

def __replace_dims__(self, dims: Tuple[str, ...], new_dims: Shape, **kwargs) -> 'Shapable':
    """
    Exchange existing dimensions.
    This can be used to rename dimensions, change dimension types or change item names.

    Args:
        dims: Dimensions to be replaced.
        new_dims: Replacement dimensions as `Shape` with `rank == len(dims)`.
        **kwargs: Additional keyword arguments required by specific implementations.
            Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
            Adding batch dimensions must always work without keyword arguments.

    Returns:
        New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
    """
    raise NotImplementedError

def __unpack_dim__(self, dim: str, unpacked_dims: phiml.math._shape.Shape, **kwargs) ‑> Shapable

Decompresses a tensor dimension by unstacking the elements along it. The compressed dimension dim is assumed to contain elements laid out according to the order of unpacked_dims.

Args

dim

Dimension to be decompressed.

unpacked_dims

Shape: Ordered dimensions to replace dim, fulfilling unpacked_dims.volume == shape(self)[dim].rank.

**kwargs

Additional keyword arguments required by specific implementations. Adding spatial dimensions to fields requires the bounds: Box argument specifying the physical extent of the new dimensions. Adding batch dimensions must always work without keyword arguments.

Returns

New instance of Shapable or NotImplemented to revert to default behavior for this object.

Expand source code

def __unpack_dim__(self, dim: str, unpacked_dims: Shape, **kwargs) -> 'Shapable':
    """
    Decompresses a tensor dimension by unstacking the elements along it.
    The compressed dimension `dim` is assumed to contain elements laid out according to the order of `unpacked_dims`.

    Args:
        dim: Dimension to be decompressed.
        unpacked_dims: `Shape`: Ordered dimensions to replace `dim`, fulfilling `unpacked_dims.volume == shape(self)[dim].rank`.
        **kwargs: Additional keyword arguments required by specific implementations.
            Adding spatial dimensions to fields requires the `bounds: Box` argument specifying the physical extent of the new dimensions.
            Adding batch dimensions must always work without keyword arguments.

    Returns:
        New instance of `Shapable` or `NotImplemented` to revert to default behavior for this object.
    """
    raise NotImplementedError

class Shaped

To be considered shaped, an object must either implement the magic method __shape__() or have a valid shape property. In either case, the returned shape must be an instance of Shape.

To check whether an object is Shaped, use isinstance(obj, Shaped).

Usage in phiml.math:

The functions shape() as well as dimension filters, such as spatial() or non_batch() can be called on all shaped objects.

See Also: Sliceable, Shapable

Expand source code

class Shaped(metaclass=_ShapedType):
    """
    To be considered shaped, an object must either implement the magic method `__shape__()` or have a valid `shape` property.
    In either case, the returned shape must be an instance of `phiml.math.Shape`.

    To check whether an object is `Shaped`, use `isinstance(obj, Shaped)`.

    **Usage in `phiml.math`:**

    The functions `phiml.math.shape` as well as dimension filters, such as `phiml.math.spatial` or `phiml.math.non_batch` can be called on all shaped objects.

    See Also:
        `Sliceable`, `Shapable`
    """

    def __shape__(self) -> 'Shape':
        """
        Returns the shape of this object.

        Alternatively, the shape can be declared via the property `shape`.

        Returns:
            `phiml.math.Shape`
        """
        raise NotImplementedError

    @property
    def shape(self) -> 'Shape':
        """
        Alternative form of `__shape__()`.
        Implement either to be considered `Shaped`.

        Returns:
            `phiml.math.Shape`
        """
        raise NotImplementedError

Instance variables

var shape : phiml.math._shape.Shape

Alternative form of __shape__(). Implement either to be considered Shaped.

Returns

Shape

Expand source code

@property
def shape(self) -> 'Shape':
    """
    Alternative form of `__shape__()`.
    Implement either to be considered `Shaped`.

    Returns:
        `phiml.math.Shape`
    """
    raise NotImplementedError

Methods

def __shape__(self) ‑> phiml.math._shape.Shape

Returns the shape of this object.

Alternatively, the shape can be declared via the property shape.

Returns

Shape

Expand source code

def __shape__(self) -> 'Shape':
    """
    Returns the shape of this object.

    Alternatively, the shape can be declared via the property `shape`.

    Returns:
        `phiml.math.Shape`
    """
    raise NotImplementedError

class Sliceable

Objects are considered sliceable if they are Shaped and implement __getitem__ as defined below.

To enable the slicing syntax obj.dim[slice], implement the __getattr__ method as defined below.

Classes implementing Sliceable should override __getattr__ to enable the special slicing syntax defined in BoundDim.

Usage in phiml.math:

In addition to slicing, sliceable objects can be unstacked along one or multiple dimensions using unstack().

See Also Shapable, Shaped

Expand source code

class Sliceable(metaclass=_SliceableType):
    """
    Objects are considered sliceable if they are `Shaped` and implement `__getitem__` as defined below.

    To enable the slicing syntax `obj.dim[slice]`, implement the `__getattr__` method as defined below.

    Classes implementing `Sliceable` should override `__getattr__` to enable the special slicing syntax defined in `BoundDim`.

    **Usage in `phiml.math`:**

    In addition to slicing, sliceable objects can be unstacked along one or multiple dimensions using `phiml.math.unstack`.

    See Also
        `Shapable`, `Shaped`
    """

    def __getitem__(self, item) -> 'Sliceable':
        """
        Slice this object along one or multiple existing or non-existing dimensions.

        When overriding this function, make sure to first call `slicing_dict(self, item)` to sort slices by dimension.

        Args:
            item: `dict` mapping dimension names to the corresponding selections.
                Selections can be slices, indices, tuples, item names, bool tensors, int tensors or other custom types.
                All Sliceable object must support indexing by `int`, `slice`, `tuple`, `list`, `str`.

        Returns:
            Instance of the same class (or a compatible class) as `self`.
        """
        raise NotImplementedError

    def __unstack__(self, dims: Tuple[str, ...]) -> Tuple['Sliceable', ...]:
        """
        Un-stack this object along one or multiple dimensions.
        Un-stacking along multiple dimensions is equal to first packing the dimensions and then unstacking along the packed dimension.

        Implementing this magic method is optional but the default implementation may be slow.

        Args:
            dims: Ordered `tuple` of dimension names along which to unstack this object.

        Returns:
            `tuple` of slices along `dims` or `NotImplemented` to revert to default behavior for this object.
        """
        raise NotImplementedError

Methods

def __getitem__(self, item) ‑> Sliceable

Slice this object along one or multiple existing or non-existing dimensions.

When overriding this function, make sure to first call slicing_dict()(self, item) to sort slices by dimension.

Args

item

dict mapping dimension names to the corresponding selections. Selections can be slices, indices, tuples, item names, bool tensors, int tensors or other custom types. All Sliceable object must support indexing by int, slice, tuple, list, str.

Returns

Instance of the same class (or a compatible class) as self.

Expand source code

def __getitem__(self, item) -> 'Sliceable':
    """
    Slice this object along one or multiple existing or non-existing dimensions.

    When overriding this function, make sure to first call `slicing_dict(self, item)` to sort slices by dimension.

    Args:
        item: `dict` mapping dimension names to the corresponding selections.
            Selections can be slices, indices, tuples, item names, bool tensors, int tensors or other custom types.
            All Sliceable object must support indexing by `int`, `slice`, `tuple`, `list`, `str`.

    Returns:
        Instance of the same class (or a compatible class) as `self`.
    """
    raise NotImplementedError

def __unstack__(self, dims: Tuple[str, ...]) ‑> Tuple[Sliceable, ...]

Un-stack this object along one or multiple dimensions. Un-stacking along multiple dimensions is equal to first packing the dimensions and then unstacking along the packed dimension.

Implementing this magic method is optional but the default implementation may be slow.

Args

dims

Ordered tuple of dimension names along which to unstack this object.

Returns

tuple of slices along dims or NotImplemented to revert to default behavior for this object.

Expand source code

def __unstack__(self, dims: Tuple[str, ...]) -> Tuple['Sliceable', ...]:
    """
    Un-stack this object along one or multiple dimensions.
    Un-stacking along multiple dimensions is equal to first packing the dimensions and then unstacking along the packed dimension.

    Implementing this magic method is optional but the default implementation may be slow.

    Args:
        dims: Ordered `tuple` of dimension names along which to unstack this object.

    Returns:
        `tuple` of slices along `dims` or `NotImplemented` to revert to default behavior for this object.
    """
    raise NotImplementedError