| 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 | |