Generic Tensor

* Support type declarations Tensor[T]
* Add type hints to wrap(), tensor(), layout(), linspace(), arange(), *range(), unstack()

Author: p-holl

phiml/math/_magic_ops.py

@@ -15,6 +15,9 @@
 from ..backend._dtype import DType
 
 
+MagicType = TypeVar('MagicType')
+OtherMagicType = TypeVar('OtherMagicType')
+
 PhiTreeNodeType = TypeVar('PhiTreeNodeType')  # Defined in phiml.math.magic: tuple, list, dict, custom
 
 
@@ -71,7 +74,7 @@ def slice_(value: PhiTreeNodeType, slices: Union[Dict[str, Union[int, slice, str
     raise ValueError(f"value must be a PhiTreeNode but got {type(value)}")
 
 
-def unstack(value, dim: DimFilter, expand=False) -> tuple:
+def unstack(value: MagicType, dim: DimFilter, expand=False) -> Tuple[MagicType, ...]:
     """
     Un-stacks a `Sliceable` along one or multiple dimensions.
 
@@ -989,9 +992,6 @@ def replace(obj: PhiTreeNodeType, **updates) -> PhiTreeNodeType:
 
 # Other Ops
 
-MagicType = TypeVar('MagicType')
-OtherMagicType = TypeVar('OtherMagicType')
-
 
 def cast(x: MagicType, dtype: Union[DType, type]) -> OtherMagicType:
     """

phiml/math/_ops.py

@@ -731,7 +731,7 @@ def meshgrid(dims: Union[Callable, Shape] = spatial, stack_dim: Union[Shape, str
     return stack_tensors(channels, stack_dim)
 
 
-def linspace(start: Union[float, Tensor, tuple, list], stop: Union[float, Tensor, tuple, list], dim: Shape) -> Tensor:
+def linspace(start: Union[float, Tensor, tuple, list], stop: Union[float, Tensor, tuple, list], dim: Shape) -> Tensor[float]:
     """
     Returns `number` evenly spaced numbers between `start` and `stop` along `dim`.
 
@@ -772,7 +772,7 @@ def linspace(start: Union[float, Tensor, tuple, list], stop: Union[float, Tensor
         return map_(linspace, start, stop, dim=dim)
 
 
-def arange(dim: Shape, start_or_stop: Union[int, None] = None, stop: Union[int, None] = None, step=1, backend=None):
+def arange(dim: Shape, start_or_stop: Union[int, None] = None, stop: Union[int, None] = None, step=1, backend=None) -> Tensor[int]:
     """
     Returns evenly spaced values between `start` and `stop`.
     If only one limit is given, `0` is used for the start.
@@ -819,7 +819,7 @@ def batched_range(dims: Shape, start: Tensor, stop: Tensor, step: Tensor):
     return batched_range(dim, start, stop, step)
 
 
-def range_tensor(*shape: Shape):
+def range_tensor(*shape: Shape) -> Tensor[int]:
     """
     Returns a `Tensor` with given `shape` containing the linear indices of each element.
     For 1D tensors, this equivalent to `arange()` with `step=1`.
@@ -838,31 +838,31 @@ def range_tensor(*shape: Shape):
     return unpack_dim(data, 'range', shape)
 
 
-def brange(start: int = 0, **stop: int):
+def brange(start: int = 0, **stop: int) -> Tensor[int]:
     """ Construct a range `Tensor` along one batch dim. """
     assert len(stop) == 1, f"brange() requires exactly one stop dimension but got {stop}"
     return arange(batch(next(iter(stop))), start, next(iter(stop.values())))
 
 
-def drange(start: int = 0, **stop: int):
+def drange(start: int = 0, **stop: int) -> Tensor[int]:
     """ Construct a range `Tensor` along one dual dim. """
     assert len(stop) == 1, f"drange() requires exactly one stop dimension but got {stop}"
     return arange(dual(next(iter(stop))), start, next(iter(stop.values())))
 
 
-def irange(start: int = 0, **stop: int):
+def irange(start: int = 0, **stop: int) -> Tensor[int]:
     """ Construct a range `Tensor` along one instance dim. """
     assert len(stop) == 1, f"irange() requires exactly one stop dimension but got {stop}"
     return arange(instance(next(iter(stop))), start, next(iter(stop.values())))
 
 
-def srange(start: int = 0, **stop: int):
+def srange(start: int = 0, **stop: int) -> Tensor[int]:
     """ Construct a range `Tensor` along one spatial dim. """
     assert len(stop) == 1, f"srange() requires exactly one stop dimension but got {stop}"
     return arange(spatial(next(iter(stop))), start, next(iter(stop.values())))
 
 
-def crange(start: int = 0, **stop: int):
+def crange(start: int = 0, **stop: int) -> Tensor[int]:
     """ Construct a range `Tensor` along one channel dim. """
     assert len(stop) == 1, f"crange() requires exactly one stop dimension but got {stop}"
     return arange(channel(next(iter(stop))), start, next(iter(stop.values())))

phiml/math/_tensors.py

@@ -5,7 +5,7 @@
 import traceback
 import warnings
 from contextlib import contextmanager
-from typing import Union, TypeVar, Sequence, Any, Literal
+from typing import Union, TypeVar, Sequence, Any, Literal, Generic
 
 from dataclasses import dataclass
 from typing import Tuple, Callable, List
@@ -28,7 +28,9 @@
 from ..backend.xops import ExtraOperator
 
 
-class Tensor:
+T = TypeVar('T')  # data type of tensors
+
+class Tensor(Generic[T]):
     """
     Abstract base class to represent structured data of one data type.
     This class replaces the native tensor classes `numpy.ndarray`, `torch.Tensor`, `tensorflow.Tensor` or `jax.numpy.ndarray` as the main data container in Φ-ML.
@@ -470,13 +472,13 @@ def _getitem(self, selection: dict) -> 'Tensor':
     def __setitem__(self, key, value):
         raise SyntaxError("Tensors are not editable to preserve the autodiff chain. This feature might be added in the future. To update part of a tensor, use math.where() or math.scatter()")
 
-    def __unstack__(self, dims: Tuple[str, ...]) -> Tuple['Tensor', ...]:  # from phiml.math.magic.Sliceable
+    def __unstack__(self, dims: Tuple[str, ...]) -> Tuple['Tensor[T]', ...]:  # from phiml.math.magic.Sliceable
         if len(dims) == 1:
             return self._unstack(dims[0])
         else:
             return NotImplemented
 
-    def _unstack(self, dim: str):
+    def _unstack(self, dim: str) -> Tuple['Tensor[T]', ...]:
         """
         Splits this tensor along the specified dimension.
         The returned tensors have the same dimensions as this tensor save the unstacked dimension.
@@ -568,12 +570,12 @@ def dimension(self, name: Union[str, Shape]) -> 'TensorDim':
         else:
             raise ValueError(name)
 
-    def pack(self, dims, packed_dim):
+    def pack(self, dims, packed_dim) -> 'Tensor[T]':
         """ See `pack_dims()` """
         from ._ops import pack_dims
         return pack_dims(self, dims, packed_dim)
 
-    def unpack(self, dim, unpacked_dims):
+    def unpack(self, dim, unpacked_dims) -> 'Tensor[T]':
         """ See `unpack_dim()` """
         from ._ops import unpack_dim
         return unpack_dim(self, dim, unpacked_dims)
@@ -584,15 +586,15 @@ def T(self):
         # return self._with_shape_replaced(self.shape.transposed())
 
     @property
-    def Ti(self):
+    def Ti(self) -> 'Tensor[T]':
         return self._with_shape_replaced(self.shape.transpose(INSTANCE_DIM))
 
     @property
-    def Tc(self):
+    def Tc(self) -> 'Tensor[T]':
         return self._with_shape_replaced(self.shape.transpose(CHANNEL_DIM))
 
     @property
-    def Ts(self):
+    def Ts(self) -> 'Tensor[T]':
         return self._with_shape_replaced(self.shape.transpose(SPATIAL_DIM))
 
     def map(self, function: Callable, dims=shape_, range=range, unwrap_scalars=True, **kwargs):
@@ -676,7 +678,7 @@ def __mod__(self, other):
     def __rmod__(self, other):
         return self._op2(other, operator.mod, True)
 
-    def __eq__(self, other) -> 'Tensor':
+    def __eq__(self, other) -> 'Tensor[bool]':
         if self is other:
             return expand(True, self.shape)
         if _EQUALITY_REDUCE[-1]['type'] == 'ref':
@@ -693,7 +695,7 @@ def __eq__(self, other) -> 'Tensor':
         else:
             return wrap(False)
 
-    def __ne__(self, other) -> 'Tensor':
+    def __ne__(self, other) -> 'Tensor[bool]':
         if _EQUALITY_REDUCE[-1]['type'] == 'ref':
             return wrap(self is not other)
         elif _EQUALITY_REDUCE[-1]['type'] == 'shape_and_value':
@@ -708,49 +710,49 @@ def __ne__(self, other) -> 'Tensor':
         else:
             return wrap(True)
 
-    def __lt__(self, other):
+    def __lt__(self, other) -> 'Tensor[bool]':
         return self._op2(other, operator.gt, True)
 
-    def __le__(self, other):
+    def __le__(self, other) -> 'Tensor[bool]':
         return self._op2(other, operator.ge, True)
 
-    def __gt__(self, other):
+    def __gt__(self, other) -> 'Tensor[bool]':
         return self._op2(other, operator.gt, False)
 
-    def __ge__(self, other):
+    def __ge__(self, other) -> 'Tensor[bool]':
         return self._op2(other, operator.ge, False)
 
-    def __lshift__(self, other):
+    def __lshift__(self, other) -> 'Tensor[T]':
         return self._op2(other, operator.lshift, False)
 
-    def __rlshift__(self, other):
+    def __rlshift__(self, other) -> 'Tensor[T]':
         return self._op2(other, operator.lshift, True)
 
-    def __rshift__(self, other):
+    def __rshift__(self, other) -> 'Tensor[T]':
         return self._op2(other, operator.rshift, False)
 
-    def __rrshift__(self, other):
+    def __rrshift__(self, other) -> 'Tensor[T]':
         return self._op2(other, operator.rshift, True)
 
-    def __abs__(self):
+    def __abs__(self) -> 'Tensor[T]':
         return self._op1(lambda t: choose_backend(t).abs(t))
 
-    def __round__(self, n=None):
+    def __round__(self, n=None) -> 'Tensor[int]':
         return self._op1(lambda t: choose_backend(t).round(t))
 
-    def __copy__(self):
+    def __copy__(self) -> 'Tensor[T]':
         return self._op1(lambda t: choose_backend(t).copy(t, only_mutable=True))
 
-    def __deepcopy__(self, memodict={}):
+    def __deepcopy__(self, memodict={}) -> 'Tensor[T]':
         return self._op1(lambda t: choose_backend(t).copy(t, only_mutable=False))
 
-    def __neg__(self) -> 'Tensor':
+    def __neg__(self) -> 'Tensor[T]':
         return self._op1(operator.neg)
 
-    def __invert__(self) -> 'Tensor':
+    def __invert__(self) -> 'Tensor[T]':
         return self._op1(lambda t: choose_backend(t).invert(t))
 
-    def __reversed__(self):
+    def __reversed__(self) -> 'Tensor[T]':
         assert self.shape.channel.rank == 1
         return self[::-1]
 
@@ -763,11 +765,11 @@ def __iter__(self):
             native = self.native([self.shape])
             return iter(native)
 
-    def item(self):
+    def item(self) -> T:
         assert self.shape.volume == 1, f"Tensor.item() is only available for single-element Tensors but got {self.shape}"
         return next(iter(self))
 
-    def __matmul__(self, other):
+    def __matmul__(self, other) -> 'Tensor[bool]':
         from ._ops import dot
         assert isinstance(other, Tensor), f"Matmul '@' requires two Tensor arguments but got {type(other)}"
         if not self.shape.dual_rank and self.shape.channel_rank:
@@ -953,7 +955,7 @@ def __init__(self, obj, stack_dim: Shape):
                 warnings.warn(f"Empty stack_dim for Layout with value {obj}")
 
     @staticmethod
-    def _recursive_get_shapes(obj, s: Shape) -> Tuple[Shape]:
+    def _recursive_get_shapes(obj, s: Shape) -> Tuple[Shape, ...]:
         if not s:
             return shape(obj, allow_unshaped=True),
         elif isinstance(obj, (tuple, list)):
@@ -1670,10 +1672,10 @@ def _simplify(self):
         return self
 
 
-def tensor(data,
+def tensor(data: Union[Sequence[T], T],
            *shape: Union[Shape, str, list],
            convert: bool = True,
-           default_list_dim=channel('vector')) -> Tensor:  # TODO assume convert_unsupported, add convert_external=False for constants
+           default_list_dim=channel('vector')) -> Tensor[T]:  # TODO assume convert_unsupported, add convert_external=False for constants
     """
     Create a Tensor from the specified `data`.
     If `convert=True`, converts `data` to the preferred format of the default backend.
@@ -1814,12 +1816,12 @@ def tensor(data,
         raise ValueError(f"{type(data)} is not supported. Only (Tensor, tuple, list, np.ndarray, native tensors) are allowed.\nCurrent backends: {BACKENDS}")
 
 
-def wrap(data, *shape: Union[Shape, str, list], default_list_dim=channel('vector')) -> Tensor:
+def wrap(data: Union[Sequence[T], T], *shape: Union[Shape, str, list], default_list_dim=channel('vector')) -> Tensor[T]:
     """ Short for `phiml.math.tensor()` with `convert=False`. """
     return tensor(data, *shape, convert=False, default_list_dim=default_list_dim)
 
 
-def layout(objects, *shape: Union[Shape, str]) -> Tensor:
+def layout(objects: Union[Sequence[T], T], *shape: Union[Shape, str]) -> Tensor[T]:
     """
     Wraps a Python tree in a `Tensor`, allowing elements to be accessed via dimensions.
     A python tree is a structure of nested `tuple`, `list`, `dict` and *leaf* objects where leaves can be any Python object.