|
|
__author__ = "Alex Rogozhnikov" |
|
|
|
|
|
from typing import Any, Dict |
|
|
|
|
|
|
|
|
from ..einops import TransformRecipe, _apply_recipe, _prepare_recipes_for_all_dims, get_backend |
|
|
from .. import EinopsError |
|
|
|
|
|
|
|
|
class RearrangeMixin: |
|
|
""" |
|
|
Rearrange layer behaves identically to einops.rearrange operation. |
|
|
|
|
|
:param pattern: str, rearrangement pattern |
|
|
:param axes_lengths: any additional specification of dimensions |
|
|
|
|
|
See einops.rearrange for source_examples. |
|
|
""" |
|
|
|
|
|
def __init__(self, pattern: str, **axes_lengths: Any) -> None: |
|
|
super().__init__() |
|
|
self.pattern = pattern |
|
|
self.axes_lengths = axes_lengths |
|
|
|
|
|
self._multirecipe = self.multirecipe() |
|
|
self._axes_lengths = tuple(self.axes_lengths.items()) |
|
|
|
|
|
def __repr__(self) -> str: |
|
|
params = repr(self.pattern) |
|
|
for axis, length in self.axes_lengths.items(): |
|
|
params += ", {}={}".format(axis, length) |
|
|
return "{}({})".format(self.__class__.__name__, params) |
|
|
|
|
|
def multirecipe(self) -> Dict[int, TransformRecipe]: |
|
|
try: |
|
|
return _prepare_recipes_for_all_dims( |
|
|
self.pattern, operation="rearrange", axes_names=tuple(self.axes_lengths) |
|
|
) |
|
|
except EinopsError as e: |
|
|
raise EinopsError(" Error while preparing {!r}\n {}".format(self, e)) |
|
|
|
|
|
def _apply_recipe(self, x): |
|
|
backend = get_backend(x) |
|
|
return _apply_recipe( |
|
|
backend=backend, |
|
|
recipe=self._multirecipe[len(x.shape)], |
|
|
tensor=x, |
|
|
reduction_type="rearrange", |
|
|
axes_lengths=self._axes_lengths, |
|
|
) |
|
|
|
|
|
def __getstate__(self): |
|
|
return {"pattern": self.pattern, "axes_lengths": self.axes_lengths} |
|
|
|
|
|
def __setstate__(self, state): |
|
|
self.__init__(pattern=state["pattern"], **state["axes_lengths"]) |
|
|
|
|
|
|
|
|
class ReduceMixin: |
|
|
""" |
|
|
Reduce layer behaves identically to einops.reduce operation. |
|
|
|
|
|
:param pattern: str, rearrangement pattern |
|
|
:param reduction: one of available reductions ('min', 'max', 'sum', 'mean', 'prod'), case-sensitive |
|
|
:param axes_lengths: any additional specification of dimensions |
|
|
|
|
|
See einops.reduce for source_examples. |
|
|
""" |
|
|
|
|
|
def __init__(self, pattern: str, reduction: str, **axes_lengths: Any): |
|
|
super().__init__() |
|
|
self.pattern = pattern |
|
|
self.reduction = reduction |
|
|
self.axes_lengths = axes_lengths |
|
|
self._multirecipe = self.multirecipe() |
|
|
self._axes_lengths = tuple(self.axes_lengths.items()) |
|
|
|
|
|
def __repr__(self): |
|
|
params = "{!r}, {!r}".format(self.pattern, self.reduction) |
|
|
for axis, length in self.axes_lengths.items(): |
|
|
params += ", {}={}".format(axis, length) |
|
|
return "{}({})".format(self.__class__.__name__, params) |
|
|
|
|
|
def multirecipe(self) -> Dict[int, TransformRecipe]: |
|
|
try: |
|
|
return _prepare_recipes_for_all_dims( |
|
|
self.pattern, operation=self.reduction, axes_names=tuple(self.axes_lengths) |
|
|
) |
|
|
except EinopsError as e: |
|
|
raise EinopsError(" Error while preparing {!r}\n {}".format(self, e)) |
|
|
|
|
|
def _apply_recipe(self, x): |
|
|
backend = get_backend(x) |
|
|
return _apply_recipe( |
|
|
backend=backend, |
|
|
recipe=self._multirecipe[len(x.shape)], |
|
|
tensor=x, |
|
|
reduction_type=self.reduction, |
|
|
axes_lengths=self._axes_lengths, |
|
|
) |
|
|
|
|
|
def __getstate__(self): |
|
|
return {"pattern": self.pattern, "reduction": self.reduction, "axes_lengths": self.axes_lengths} |
|
|
|
|
|
def __setstate__(self, state): |
|
|
self.__init__(pattern=state["pattern"], reduction=state["reduction"], **state["axes_lengths"]) |
|
|
|
