| from typing import List, Optional, Dict, cast |
|
|
| import tensorflow as tf |
| from tensorflow.keras.layers import Layer |
|
|
| from .._backends import UnknownSize |
| from . import RearrangeMixin, ReduceMixin |
| from ._einmix import _EinmixMixin |
| from ..einops import TransformRecipe, _reconstruct_from_shape_uncached |
|
|
| __author__ = 'Alex Rogozhnikov' |
|
|
|
|
| def _compute_output_shape(recipe: TransformRecipe, input_shape) -> List[Optional[int]]: |
| input_shape = [UnknownSize() if d is None else int(d) for d in input_shape] |
| init_shapes, reduced_axes, axes_reordering, added_axes, final_shape = \ |
| _reconstruct_from_shape_uncached(recipe, input_shape) |
| output_shape: List[Optional[int]] = [None if isinstance(d, UnknownSize) else int(d) for d in final_shape] |
| return output_shape |
|
|
|
|
| class Rearrange(RearrangeMixin, Layer): |
| def compute_output_shape(self, input_shape): |
| return _compute_output_shape(self.recipe(), input_shape) |
|
|
| def call(self, inputs): |
| return self._apply_recipe(inputs) |
|
|
| def get_config(self): |
| return {'pattern': self.pattern, **self.axes_lengths} |
|
|
|
|
| class Reduce(ReduceMixin, Layer): |
| def compute_output_shape(self, input_shape): |
| return _compute_output_shape(self.recipe(), input_shape) |
|
|
| def call(self, inputs): |
| return self._apply_recipe(inputs) |
|
|
| def get_config(self): |
| return {'pattern': self.pattern, 'reduction': self.reduction, **self.axes_lengths} |
|
|
|
|
| class EinMix(_EinmixMixin, Layer): |
| def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound): |
| self.weight = tf.Variable(tf.random_uniform_initializer(-weight_bound, weight_bound)(shape=weight_shape), |
| trainable=True) |
| if bias_shape is not None: |
| self.bias = tf.Variable(tf.random_uniform_initializer(-bias_bound, bias_bound)(shape=bias_shape), |
| trainable=True) |
| else: |
| self.bias = None |
|
|
| def _create_rearrange_layers(self, |
| pre_reshape_pattern: Optional[str], |
| pre_reshape_lengths: Optional[Dict], |
| post_reshape_pattern: Optional[str], |
| post_reshape_lengths: Optional[Dict], |
| ): |
| self.pre_rearrange = None |
| if pre_reshape_pattern is not None: |
| self.pre_rearrange = Rearrange(pre_reshape_pattern, **cast(dict, pre_reshape_lengths)) |
|
|
| self.post_rearrange = None |
| if post_reshape_pattern is not None: |
| self.post_rearrange = Rearrange(post_reshape_pattern, **cast(dict, post_reshape_lengths)) |
|
|
| def build(self, input_shape): |
| pass |
|
|
| def call(self, inputs): |
| if self.pre_rearrange is not None: |
| inputs = self.pre_rearrange(inputs) |
| result = tf.einsum(self.einsum_pattern, inputs, self.weight) |
| if self.bias is not None: |
| result = result + self.bias |
| if self.post_rearrange is not None: |
| result = self.post_rearrange(result) |
| return result |
|
|
| def get_config(self): |
| return {'pattern': self.pattern, |
| 'weight_shape': self.weight_shape, |
| 'bias_shape': self.bias_shape, |
| **self.axes_lengths} |
|
|
