Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,105 @@
|
||||
__author__ = "Alex Rogozhnikov"
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from einops import EinopsError
|
||||
from einops.einops import TransformRecipe, _apply_recipe, _prepare_recipes_for_all_dims, get_backend
|
||||
|
||||
|
||||
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._recipe = self.recipe() # checking parameters
|
||||
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 += f", {axis}={length}"
|
||||
return f"{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(f" Error while preparing {self!r}\n {e}") from None
|
||||
|
||||
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 = f"{self.pattern!r}, {self.reduction!r}"
|
||||
for axis, length in self.axes_lengths.items():
|
||||
params += f", {axis}={length}"
|
||||
return f"{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(f" Error while preparing {self!r}\n {e}") from None
|
||||
|
||||
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"])
|
||||
@@ -0,0 +1,227 @@
|
||||
import string
|
||||
import warnings
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from einops import EinopsError
|
||||
from einops.einops import _product
|
||||
from einops.parsing import ParsedExpression, _ellipsis
|
||||
|
||||
|
||||
def _report_axes(axes: set, report_message: str):
|
||||
if len(axes) > 0:
|
||||
raise EinopsError(report_message.format(axes))
|
||||
|
||||
|
||||
class _EinmixMixin:
|
||||
def __init__(self, pattern: str, weight_shape: str, bias_shape: Optional[str] = None, **axes_lengths: Any):
|
||||
"""
|
||||
EinMix - Einstein summation with automated tensor management and axis packing/unpacking.
|
||||
|
||||
EinMix is a combination of einops and MLP, see tutorial:
|
||||
https://github.com/arogozhnikov/einops/blob/main/docs/3-einmix-layer.ipynb
|
||||
|
||||
Imagine taking einsum with two arguments, one of each input, and one - tensor with weights
|
||||
>>> einsum('time batch channel_in, channel_in channel_out -> time batch channel_out', input, weight)
|
||||
|
||||
This layer manages weights for you, syntax highlights a special role of weight matrix
|
||||
>>> EinMix('time batch channel_in -> time batch channel_out', weight_shape='channel_in channel_out')
|
||||
But otherwise it is the same einsum under the hood. Plus einops-rearrange.
|
||||
|
||||
Simple linear layer with a bias term (you have one like that in your framework)
|
||||
>>> EinMix('t b cin -> t b cout', weight_shape='cin cout', bias_shape='cout', cin=10, cout=20)
|
||||
There is no restriction to mix the last axis. Let's mix along height
|
||||
>>> EinMix('h w c-> hout w c', weight_shape='h hout', bias_shape='hout', h=32, hout=32)
|
||||
Example of channel-wise multiplication (like one used in normalizations)
|
||||
>>> EinMix('t b c -> t b c', weight_shape='c', c=128)
|
||||
Multi-head linear layer (each head is own linear layer):
|
||||
>>> EinMix('t b (head cin) -> t b (head cout)', weight_shape='head cin cout', ...)
|
||||
|
||||
... and yes, you need to specify all dimensions of weight shape/bias shape in parameters.
|
||||
|
||||
Use cases:
|
||||
- when channel dimension is not last, use EinMix, not transposition
|
||||
- patch/segment embeddings
|
||||
- when need only within-group connections to reduce number of weights and computations
|
||||
- next-gen MLPs (follow tutorial link above to learn more!)
|
||||
- in general, any time you want to combine linear layer and einops.rearrange
|
||||
|
||||
Uniform He initialization is applied to weight tensor.
|
||||
This accounts for the number of elements mixed and produced.
|
||||
|
||||
Parameters
|
||||
:param pattern: transformation pattern, left side - dimensions of input, right side - dimensions of output
|
||||
:param weight_shape: axes of weight. A tensor of this shape is created, stored, and optimized in a layer
|
||||
If bias_shape is not specified, bias is not created.
|
||||
:param bias_shape: axes of bias added to output. Weights of this shape are created and stored. If `None` (the default), no bias is added.
|
||||
:param axes_lengths: dimensions of weight tensor
|
||||
"""
|
||||
super().__init__()
|
||||
self.pattern = pattern
|
||||
self.weight_shape = weight_shape
|
||||
self.bias_shape = bias_shape
|
||||
self.axes_lengths = axes_lengths
|
||||
self.initialize_einmix(
|
||||
pattern=pattern, weight_shape=weight_shape, bias_shape=bias_shape, axes_lengths=axes_lengths
|
||||
)
|
||||
|
||||
def initialize_einmix(self, pattern: str, weight_shape: str, bias_shape: Optional[str], axes_lengths: dict):
|
||||
left_pattern, right_pattern = pattern.split("->")
|
||||
left = ParsedExpression(left_pattern)
|
||||
right = ParsedExpression(right_pattern)
|
||||
weight = ParsedExpression(weight_shape)
|
||||
_report_axes(
|
||||
set.difference(right.identifiers, {*left.identifiers, *weight.identifiers}),
|
||||
"Unrecognized identifiers on the right side of EinMix {}",
|
||||
)
|
||||
if weight.has_ellipsis:
|
||||
raise EinopsError("Ellipsis is not supported in weight, as its shape should be fully specified")
|
||||
if left.has_ellipsis or right.has_ellipsis:
|
||||
if not (left.has_ellipsis and right.has_ellipsis):
|
||||
raise EinopsError(f"Ellipsis in EinMix should be on both sides, {pattern}")
|
||||
if left.has_ellipsis_parenthesized:
|
||||
raise EinopsError(f"Ellipsis on left side can't be in parenthesis, got {pattern}")
|
||||
if any(x.has_non_unitary_anonymous_axes for x in [left, right, weight]):
|
||||
raise EinopsError("Anonymous axes (numbers) are not allowed in EinMix")
|
||||
if "(" in weight_shape or ")" in weight_shape:
|
||||
raise EinopsError(f"Parenthesis is not allowed in weight shape: {weight_shape}")
|
||||
|
||||
pre_reshape_pattern = None
|
||||
pre_reshape_lengths = None
|
||||
post_reshape_pattern = None
|
||||
if any(len(group) != 1 for group in left.composition):
|
||||
names: List[str] = []
|
||||
for group in left.composition:
|
||||
names += group
|
||||
names = [name if name != _ellipsis else "..." for name in names]
|
||||
composition = " ".join(names)
|
||||
pre_reshape_pattern = f"{left_pattern}-> {composition}"
|
||||
pre_reshape_lengths = {name: length for name, length in axes_lengths.items() if name in names}
|
||||
|
||||
if any(len(group) != 1 for group in right.composition) or right.has_ellipsis_parenthesized:
|
||||
names = []
|
||||
for group in right.composition:
|
||||
names += group
|
||||
names = [name if name != _ellipsis else "..." for name in names]
|
||||
composition = " ".join(names)
|
||||
post_reshape_pattern = f"{composition} ->{right_pattern}"
|
||||
|
||||
self._create_rearrange_layers(pre_reshape_pattern, pre_reshape_lengths, post_reshape_pattern, {})
|
||||
|
||||
for axis in weight.identifiers:
|
||||
if axis not in axes_lengths:
|
||||
raise EinopsError(f"Dimension {axis} of weight should be specified")
|
||||
_report_axes(
|
||||
set.difference(set(axes_lengths), {*left.identifiers, *weight.identifiers}),
|
||||
"Axes {} are not used in pattern",
|
||||
)
|
||||
_report_axes(
|
||||
set.difference(weight.identifiers, {*left.identifiers, *right.identifiers}), "Weight axes {} are redundant"
|
||||
)
|
||||
if len(weight.identifiers) == 0:
|
||||
warnings.warn("EinMix: weight has no dimensions (means multiplication by a number)", stacklevel=2)
|
||||
|
||||
_weight_shape = [axes_lengths[axis] for (axis,) in weight.composition]
|
||||
# single output element is a combination of fan_in input elements
|
||||
_fan_in = _product([axes_lengths[axis] for (axis,) in weight.composition if axis not in right.identifiers])
|
||||
if bias_shape is not None:
|
||||
# maybe I should put ellipsis in the beginning for simplicity?
|
||||
if not isinstance(bias_shape, str):
|
||||
raise EinopsError("bias shape should be string specifying which axes bias depends on")
|
||||
bias = ParsedExpression(bias_shape)
|
||||
_report_axes(
|
||||
set.difference(bias.identifiers, right.identifiers),
|
||||
"Bias axes {} not present in output",
|
||||
)
|
||||
_report_axes(
|
||||
set.difference(bias.identifiers, set(axes_lengths)),
|
||||
"Sizes not provided for bias axes {}",
|
||||
)
|
||||
|
||||
_bias_shape = []
|
||||
used_non_trivial_size = False
|
||||
for axes in right.composition:
|
||||
if axes == _ellipsis:
|
||||
if used_non_trivial_size:
|
||||
raise EinopsError("all bias dimensions should go after ellipsis in the output")
|
||||
else:
|
||||
# handles ellipsis correctly
|
||||
for axis in axes:
|
||||
if axis == _ellipsis:
|
||||
if used_non_trivial_size:
|
||||
raise EinopsError("all bias dimensions should go after ellipsis in the output")
|
||||
elif axis in bias.identifiers:
|
||||
_bias_shape.append(axes_lengths[axis])
|
||||
used_non_trivial_size = True
|
||||
else:
|
||||
_bias_shape.append(1)
|
||||
else:
|
||||
_bias_shape = None
|
||||
|
||||
weight_bound = (3 / _fan_in) ** 0.5
|
||||
bias_bound = (1 / _fan_in) ** 0.5
|
||||
self._create_parameters(_weight_shape, weight_bound, _bias_shape, bias_bound)
|
||||
|
||||
# rewrite einsum expression with single-letter latin identifiers so that
|
||||
# expression will be understood by any framework
|
||||
mapped_identifiers = {*left.identifiers, *right.identifiers, *weight.identifiers}
|
||||
if _ellipsis in mapped_identifiers:
|
||||
mapped_identifiers.remove(_ellipsis)
|
||||
mapped_identifiers = sorted(mapped_identifiers)
|
||||
mapping2letters = {k: letter for letter, k in zip(string.ascii_lowercase, mapped_identifiers)}
|
||||
mapping2letters[_ellipsis] = "..." # preserve ellipsis
|
||||
|
||||
def write_flat_remapped(axes: ParsedExpression):
|
||||
result = []
|
||||
for composed_axis in axes.composition:
|
||||
if isinstance(composed_axis, list):
|
||||
result.extend([mapping2letters[axis] for axis in composed_axis])
|
||||
else:
|
||||
assert composed_axis == _ellipsis
|
||||
result.append("...")
|
||||
return "".join(result)
|
||||
|
||||
self.einsum_pattern: str = (
|
||||
f"{write_flat_remapped(left)},{write_flat_remapped(weight)}->{write_flat_remapped(right)}"
|
||||
)
|
||||
|
||||
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],
|
||||
):
|
||||
raise NotImplementedError("Should be defined in framework implementations")
|
||||
|
||||
def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
|
||||
"""Shape and implementations"""
|
||||
raise NotImplementedError("Should be defined in framework implementations")
|
||||
|
||||
def __repr__(self):
|
||||
params = repr(self.pattern)
|
||||
params += f", '{self.weight_shape}'"
|
||||
if self.bias_shape is not None:
|
||||
params += f", '{self.bias_shape}'"
|
||||
for axis, length in self.axes_lengths.items():
|
||||
params += f", {axis}={length}"
|
||||
return f"{self.__class__.__name__}({params})"
|
||||
|
||||
|
||||
class _EinmixDebugger(_EinmixMixin):
|
||||
"""Used only to test mixin"""
|
||||
|
||||
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_reshape_pattern = pre_reshape_pattern
|
||||
self.pre_reshape_lengths = pre_reshape_lengths
|
||||
self.post_reshape_pattern = post_reshape_pattern
|
||||
self.post_reshape_lengths = post_reshape_lengths
|
||||
|
||||
def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
|
||||
self.saved_weight_shape = weight_shape
|
||||
self.saved_bias_shape = bias_shape
|
||||
@@ -0,0 +1,82 @@
|
||||
from dataclasses import field
|
||||
from typing import Dict, Optional, cast
|
||||
|
||||
import flax.linen as nn
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
|
||||
from . import RearrangeMixin, ReduceMixin
|
||||
from ._einmix import _EinmixMixin
|
||||
|
||||
__author__ = "Alex Rogozhnikov"
|
||||
|
||||
|
||||
class Reduce(nn.Module):
|
||||
pattern: str
|
||||
reduction: str
|
||||
sizes: dict = field(default_factory=dict)
|
||||
|
||||
def setup(self):
|
||||
self.reducer = ReduceMixin(self.pattern, self.reduction, **self.sizes)
|
||||
|
||||
def __call__(self, input):
|
||||
return self.reducer._apply_recipe(input)
|
||||
|
||||
|
||||
class Rearrange(nn.Module):
|
||||
pattern: str
|
||||
sizes: dict = field(default_factory=dict)
|
||||
|
||||
def setup(self):
|
||||
self.rearranger = RearrangeMixin(self.pattern, **self.sizes)
|
||||
|
||||
def __call__(self, input):
|
||||
return self.rearranger._apply_recipe(input)
|
||||
|
||||
|
||||
class EinMix(nn.Module, _EinmixMixin):
|
||||
pattern: str
|
||||
weight_shape: str
|
||||
bias_shape: Optional[str] = None
|
||||
sizes: dict = field(default_factory=dict)
|
||||
|
||||
def setup(self):
|
||||
self.initialize_einmix(
|
||||
pattern=self.pattern,
|
||||
weight_shape=self.weight_shape,
|
||||
bias_shape=self.bias_shape,
|
||||
axes_lengths=self.sizes,
|
||||
)
|
||||
|
||||
def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
|
||||
self.weight = self.param("weight", jax.nn.initializers.uniform(weight_bound), weight_shape)
|
||||
|
||||
if bias_shape is not None:
|
||||
self.bias = self.param("bias", jax.nn.initializers.uniform(bias_bound), bias_shape)
|
||||
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, sizes=cast(dict, pre_reshape_lengths))
|
||||
|
||||
self.post_rearrange = None
|
||||
if post_reshape_pattern is not None:
|
||||
self.post_rearrange = Rearrange(post_reshape_pattern, sizes=cast(dict, post_reshape_lengths))
|
||||
|
||||
def __call__(self, input):
|
||||
if self.pre_rearrange is not None:
|
||||
input = self.pre_rearrange(input)
|
||||
result = jnp.einsum(self.einsum_pattern, input, self.weight)
|
||||
if self.bias is not None:
|
||||
result += self.bias
|
||||
if self.post_rearrange is not None:
|
||||
result = self.post_rearrange(result)
|
||||
return result
|
||||
@@ -0,0 +1,9 @@
|
||||
__author__ = "Alex Rogozhnikov"
|
||||
|
||||
from einops.layers.tensorflow import EinMix, Rearrange, Reduce
|
||||
|
||||
keras_custom_objects = {
|
||||
Rearrange.__name__: Rearrange,
|
||||
Reduce.__name__: Reduce,
|
||||
EinMix.__name__: EinMix,
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
from typing import Dict, Optional, cast
|
||||
|
||||
import oneflow as flow
|
||||
|
||||
from . import RearrangeMixin, ReduceMixin
|
||||
from ._einmix import _EinmixMixin
|
||||
|
||||
__author__ = "Tianhe Ren & Depeng Liang"
|
||||
|
||||
|
||||
class Rearrange(RearrangeMixin, flow.nn.Module):
|
||||
def forward(self, input):
|
||||
return self._apply_recipe(input)
|
||||
|
||||
|
||||
class Reduce(ReduceMixin, flow.nn.Module):
|
||||
def forward(self, input):
|
||||
return self._apply_recipe(input)
|
||||
|
||||
|
||||
class EinMix(_EinmixMixin, flow.nn.Module):
|
||||
def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
|
||||
self.weight = flow.nn.Parameter(
|
||||
flow.zeros(weight_shape).uniform_(-weight_bound, weight_bound), requires_grad=True
|
||||
)
|
||||
if bias_shape is not None:
|
||||
self.bias = flow.nn.Parameter(flow.zeros(bias_shape).uniform_(-bias_bound, bias_bound), requires_grad=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 forward(self, input):
|
||||
if self.pre_rearrange is not None:
|
||||
input = self.pre_rearrange(input)
|
||||
result = flow.einsum(self.einsum_pattern, input, self.weight)
|
||||
if self.bias is not None:
|
||||
result += self.bias
|
||||
if self.post_rearrange is not None:
|
||||
result = self.post_rearrange(result)
|
||||
return result
|
||||
@@ -0,0 +1,58 @@
|
||||
from typing import Dict, Optional, cast
|
||||
|
||||
import paddle
|
||||
|
||||
from . import RearrangeMixin, ReduceMixin
|
||||
from ._einmix import _EinmixMixin
|
||||
|
||||
__author__ = "PaddlePaddle"
|
||||
|
||||
|
||||
class Rearrange(RearrangeMixin, paddle.nn.Layer):
|
||||
def forward(self, input):
|
||||
return self._apply_recipe(input)
|
||||
|
||||
|
||||
class Reduce(ReduceMixin, paddle.nn.Layer):
|
||||
def forward(self, input):
|
||||
return self._apply_recipe(input)
|
||||
|
||||
|
||||
class EinMix(_EinmixMixin, paddle.nn.Layer):
|
||||
def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
|
||||
self.weight = self.create_parameter(
|
||||
weight_shape, default_initializer=paddle.nn.initializer.Uniform(-weight_bound, weight_bound)
|
||||
)
|
||||
|
||||
if bias_shape is not None:
|
||||
self.bias = self.create_parameter(
|
||||
bias_shape, default_initializer=paddle.nn.initializer.Uniform(-bias_bound, bias_bound)
|
||||
)
|
||||
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 forward(self, input):
|
||||
if self.pre_rearrange is not None:
|
||||
input = self.pre_rearrange(input)
|
||||
|
||||
result = paddle.einsum(self.einsum_pattern, input, self.weight)
|
||||
if self.bias is not None:
|
||||
result += self.bias
|
||||
if self.post_rearrange is not None:
|
||||
result = self.post_rearrange(result)
|
||||
return result
|
||||
@@ -0,0 +1,101 @@
|
||||
"""
|
||||
Comment about tensorflow layers:
|
||||
unfortunately instructions on creation of TF layers change constantly,
|
||||
and changed way too many times at this point to remember what-compatible-where.
|
||||
|
||||
Layers in einops==0.7.0 (and several prior versions)
|
||||
are compatible with TF 2.13
|
||||
|
||||
Layers in einops==0.8.0 were re-implemented
|
||||
according to official instructions for TF 2.16
|
||||
|
||||
"""
|
||||
|
||||
from typing import Dict, Optional, cast
|
||||
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras.layers import Layer
|
||||
|
||||
from . import RearrangeMixin, ReduceMixin
|
||||
from ._einmix import _EinmixMixin
|
||||
|
||||
__author__ = "Alex Rogozhnikov"
|
||||
|
||||
|
||||
class Rearrange(RearrangeMixin, Layer):
|
||||
def build(self, input_shape):
|
||||
pass # layer does not have any parameters to be initialized
|
||||
|
||||
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 build(self, input_shape):
|
||||
pass # layer does not have any parameters to be initialized
|
||||
|
||||
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):
|
||||
# this method is called in __init__,
|
||||
# but we postpone actual creation to build(), as TF instruction suggests
|
||||
self._params = [weight_shape, weight_bound, bias_shape, bias_bound]
|
||||
|
||||
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):
|
||||
[weight_shape, weight_bound, bias_shape, bias_bound] = self._params
|
||||
self.weight = self.add_weight(
|
||||
shape=weight_shape,
|
||||
initializer=tf.random_uniform_initializer(-weight_bound, weight_bound),
|
||||
trainable=True,
|
||||
)
|
||||
|
||||
if bias_shape is not None:
|
||||
self.bias = self.add_weight(
|
||||
shape=bias_shape,
|
||||
initializer=tf.random_uniform_initializer(-bias_bound, bias_bound),
|
||||
trainable=True,
|
||||
)
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
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,
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
from typing import Dict, Optional, cast
|
||||
|
||||
import torch
|
||||
|
||||
from einops._torch_specific import apply_for_scriptable_torch
|
||||
|
||||
from . import RearrangeMixin, ReduceMixin
|
||||
from ._einmix import _EinmixMixin
|
||||
|
||||
__author__ = "Alex Rogozhnikov"
|
||||
|
||||
|
||||
class Rearrange(RearrangeMixin, torch.nn.Module):
|
||||
def forward(self, input):
|
||||
recipe = self._multirecipe[input.ndim]
|
||||
return apply_for_scriptable_torch(recipe, input, reduction_type="rearrange", axes_dims=self._axes_lengths)
|
||||
|
||||
def _apply_recipe(self, x):
|
||||
# overriding parent method to prevent it's scripting
|
||||
pass
|
||||
|
||||
|
||||
class Reduce(ReduceMixin, torch.nn.Module):
|
||||
def forward(self, input):
|
||||
recipe = self._multirecipe[input.ndim]
|
||||
return apply_for_scriptable_torch(recipe, input, reduction_type=self.reduction, axes_dims=self._axes_lengths)
|
||||
|
||||
def _apply_recipe(self, x):
|
||||
# overriding parent method to prevent it's scripting
|
||||
pass
|
||||
|
||||
|
||||
class EinMix(_EinmixMixin, torch.nn.Module):
|
||||
def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
|
||||
self.weight = torch.nn.Parameter(
|
||||
torch.zeros(weight_shape).uniform_(-weight_bound, weight_bound), requires_grad=True
|
||||
)
|
||||
if bias_shape is not None:
|
||||
self.bias = torch.nn.Parameter(
|
||||
torch.zeros(bias_shape).uniform_(-bias_bound, bias_bound), requires_grad=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 forward(self, input):
|
||||
if self.pre_rearrange is not None:
|
||||
input = self.pre_rearrange(input)
|
||||
result = torch.einsum(self.einsum_pattern, input, self.weight)
|
||||
if self.bias is not None:
|
||||
result += self.bias
|
||||
if self.post_rearrange is not None:
|
||||
result = self.post_rearrange(result)
|
||||
return result
|
||||
Reference in New Issue
Block a user