83 lines
2.5 KiB
Python
83 lines
2.5 KiB
Python
|
from dataclasses import field
|
||
|
|
from typing import Optional, Dict, 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=lambda: {})
|
||
|
|
|
||
|
|
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=lambda: {})
|
||
|
|
|
||
|
|
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=lambda: {})
|
||
|
|
|
||
|
|
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
|