|
|
""" |
|
|
Specialization of einops for torch. |
|
|
|
|
|
Unfortunately, torch's jit scripting mechanism isn't strong enough, |
|
|
and to have scripting supported at least for layers, |
|
|
a number of additional moves is needed. |
|
|
|
|
|
Design of main operations (dynamic resolution by lookup) is unlikely |
|
|
to be implemented by torch.jit.script, |
|
|
but torch.compile seems to work with operations just fine. |
|
|
""" |
|
|
|
|
|
import warnings |
|
|
from typing import Dict, List, Tuple |
|
|
|
|
|
import torch |
|
|
from einops.einops import TransformRecipe, _reconstruct_from_shape_uncached |
|
|
|
|
|
|
|
|
class TorchJitBackend: |
|
|
""" |
|
|
Completely static backend that mimics part of normal backend functionality |
|
|
but restricted to be within torchscript. |
|
|
""" |
|
|
|
|
|
@staticmethod |
|
|
def reduce(x: torch.Tensor, operation: str, reduced_axes: List[int]): |
|
|
if operation == "min": |
|
|
return x.amin(dim=reduced_axes) |
|
|
elif operation == "max": |
|
|
return x.amax(dim=reduced_axes) |
|
|
elif operation == "sum": |
|
|
return x.sum(dim=reduced_axes) |
|
|
elif operation == "mean": |
|
|
return x.mean(dim=reduced_axes) |
|
|
elif operation == "prod": |
|
|
for i in list(sorted(reduced_axes))[::-1]: |
|
|
x = x.prod(dim=i) |
|
|
return x |
|
|
else: |
|
|
raise NotImplementedError("Unknown reduction ", operation) |
|
|
|
|
|
@staticmethod |
|
|
def transpose(x, axes: List[int]): |
|
|
return x.permute(axes) |
|
|
|
|
|
@staticmethod |
|
|
def stack_on_zeroth_dimension(tensors: List[torch.Tensor]): |
|
|
return torch.stack(tensors) |
|
|
|
|
|
@staticmethod |
|
|
def tile(x, repeats: List[int]): |
|
|
return x.repeat(repeats) |
|
|
|
|
|
@staticmethod |
|
|
def add_axes(x, n_axes: int, pos2len: Dict[int, int]): |
|
|
repeats = [-1] * n_axes |
|
|
for axis_position, axis_length in pos2len.items(): |
|
|
x = torch.unsqueeze(x, axis_position) |
|
|
repeats[axis_position] = axis_length |
|
|
return x.expand(repeats) |
|
|
|
|
|
@staticmethod |
|
|
def is_float_type(x): |
|
|
return x.dtype in [torch.float16, torch.float32, torch.float64, torch.bfloat16] |
|
|
|
|
|
@staticmethod |
|
|
def shape(x): |
|
|
return x.shape |
|
|
|
|
|
@staticmethod |
|
|
def reshape(x, shape: List[int]): |
|
|
return x.reshape(shape) |
|
|
|
|
|
|
|
|
|
|
|
def apply_for_scriptable_torch( |
|
|
recipe: TransformRecipe, tensor: torch.Tensor, reduction_type: str, axes_dims: List[Tuple[str, int]] |
|
|
) -> torch.Tensor: |
|
|
backend = TorchJitBackend |
|
|
( |
|
|
init_shapes, |
|
|
axes_reordering, |
|
|
reduced_axes, |
|
|
added_axes, |
|
|
final_shapes, |
|
|
n_axes_w_added, |
|
|
) = _reconstruct_from_shape_uncached(recipe, backend.shape(tensor), axes_dims=axes_dims) |
|
|
if init_shapes is not None: |
|
|
tensor = backend.reshape(tensor, init_shapes) |
|
|
if axes_reordering is not None: |
|
|
tensor = backend.transpose(tensor, axes_reordering) |
|
|
if len(reduced_axes) > 0: |
|
|
tensor = backend.reduce(tensor, operation=reduction_type, reduced_axes=reduced_axes) |
|
|
if len(added_axes) > 0: |
|
|
tensor = backend.add_axes(tensor, n_axes=n_axes_w_added, pos2len=added_axes) |
|
|
if final_shapes is not None: |
|
|
tensor = backend.reshape(tensor, final_shapes) |
|
|
return tensor |
|
|
|
|
|
|
|
|
def allow_ops_in_compiled_graph(): |
|
|
if hasattr(torch, "__version__") and torch.__version__[0] < "2": |
|
|
|
|
|
return |
|
|
try: |
|
|
from torch._dynamo import allow_in_graph |
|
|
except ImportError: |
|
|
warnings.warn("allow_ops_in_compiled_graph failed to import torch: ensure pytorch >=2.0", ImportWarning) |
|
|
return |
|
|
|
|
|
from .einops import rearrange, reduce, repeat, einsum |
|
|
from .packing import pack, unpack |
|
|
|
|
|
allow_in_graph(rearrange) |
|
|
allow_in_graph(reduce) |
|
|
allow_in_graph(repeat) |
|
|
allow_in_graph(einsum) |
|
|
allow_in_graph(pack) |
|
|
allow_in_graph(unpack) |
|
|
|
|
|
|
|
|
global _ops_were_registered_in_torchdynamo |
|
|
_ops_were_registered_in_torchdynamo = True |
|
|
|
|
|
|
|
|
|
|
|
allow_ops_in_compiled_graph() |
|
|
|
