# Path: einops/layers/tensorflow.py
# 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, **pre_reshape_lengths)
# 
#         self.post_rearrange = None
#         if post_reshape_pattern is not None:
#             self.post_rearrange = Rearrange(post_reshape_pattern, **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}


from ..layers.tensorflow import Rearrange, Reduce, EinMix
__author__ = 'Alex Rogozhnikov'


keras_custom_objects = {
    Rearrange.__name__: Rearrange,
    Reduce.__name__: Reduce,

# Path: einops/layers/tensorflow.py
# 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, **pre_reshape_lengths)
# 
#         self.post_rearrange = None
#         if post_reshape_pattern is not None:
#             self.post_rearrange = Rearrange(post_reshape_pattern, **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}


# Path: einops/layers/keras.py
from ..layers.tensorflow import Rearrange, Reduce, EinMix

__author__ = 'Alex Rogozhnikov'


keras_custom_objects = {
    Rearrange.__name__: Rearrange,
    Reduce.__name__: Reduce,

    EinMix.__name__: EinMix,