Kernels:

kernels-community
/

einops

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on about 18 hours ago

Commit

ae022fa

verified ·

1 Parent(s): cb47c7f

Uploaded using `kernel-builder`.

Browse files

Files changed (20) hide show

build/torch-cuda/__init__.py +18 -0
build/torch-cuda/_backends.py +719 -0
build/torch-cuda/_ops.py +38 -0
build/torch-cuda/_torch_specific.py +128 -0
build/torch-cuda/array_api.py +124 -0
build/torch-cuda/einops.py +916 -0
build/torch-cuda/experimental/__init__.py +0 -0
build/torch-cuda/experimental/indexing.py +5 -0
build/torch-cuda/layers/__init__.py +106 -0
build/torch-cuda/layers/_einmix.py +229 -0
build/torch-cuda/layers/flax.py +82 -0
build/torch-cuda/layers/keras.py +9 -0
build/torch-cuda/layers/oneflow.py +54 -0
build/torch-cuda/layers/paddle.py +58 -0
build/torch-cuda/layers/tensorflow.py +103 -0
build/torch-cuda/layers/torch.py +67 -0
build/torch-cuda/metadata.json +45 -0
build/torch-cuda/metadata.json.sigstore +1 -0
build/torch-cuda/packing.py +190 -0
build/torch-cuda/parsing.py +152 -0

build/torch-cuda/__init__.py ADDED Viewed

	@@ -0,0 +1,18 @@

+# imports can use EinopsError class
+# ruff: noqa: E402
+__author__ = "Alex Rogozhnikov"
+__version__ = "0.8.1"
+class EinopsError(RuntimeError):
+    """Runtime error thrown by einops"""
+    pass
+__all__ = ["rearrange", "reduce", "repeat", "einsum", "pack", "unpack", "parse_shape", "asnumpy", "EinopsError"]
+from .einops import rearrange, reduce, repeat, einsum, parse_shape, asnumpy
+from .packing import pack, unpack
+from . import array_api

build/torch-cuda/_backends.py ADDED Viewed

	@@ -0,0 +1,719 @@

+"""
+Backends in `einops` are organized to meet the following requirements
+- backends are not imported unless those are actually needed, because
+    - backends may not be installed
+    - importing all available backends will drive to significant memory footprint
+    - backends may be present but installed with errors (but never used),
+      importing may drive to crashes
+- backend should be either symbolic or imperative
+    - this determines which methods (from_numpy/to_numpy or create_symbol/eval_symbol) should be defined
+- if backend can't provide symbols for shape dimensions, UnknownSize objects are used
+"""
+import sys
+__author__ = "Alex Rogozhnikov"
+_loaded_backends: dict = {}
+_type2backend: dict = {}
+_debug_importing = False
+def get_backend(tensor) -> "AbstractBackend":
+    """
+    Takes a correct backend (e.g. numpy backend if tensor is numpy.ndarray) for a tensor.
+    If needed, imports package and creates backend
+    """
+    _type = type(tensor)
+    _result = _type2backend.get(_type, None)
+    if _result is not None:
+        return _result
+    for framework_name, backend in list(_loaded_backends.items()):
+        if backend.is_appropriate_type(tensor):
+            _type2backend[_type] = backend
+            return backend
+    # Find backend subclasses recursively
+    backend_subclasses = []
+    backends = AbstractBackend.__subclasses__()
+    while backends:
+        backend = backends.pop()
+        backends += backend.__subclasses__()
+        backend_subclasses.append(backend)
+    for BackendSubclass in backend_subclasses:
+        if _debug_importing:
+            print("Testing for subclass of ", BackendSubclass)
+        if BackendSubclass.framework_name not in _loaded_backends:
+            # check that module was already imported. Otherwise it can't be imported
+            if BackendSubclass.framework_name in sys.modules:
+                if _debug_importing:
+                    print("Imported backend for ", BackendSubclass.framework_name)
+                backend = BackendSubclass()
+                _loaded_backends[backend.framework_name] = backend
+                if backend.is_appropriate_type(tensor):
+                    _type2backend[_type] = backend
+                    return backend
+    raise RuntimeError("Tensor type unknown to einops {}".format(type(tensor)))
+class AbstractBackend:
+    """Base backend class, major part of methods are only for debugging purposes."""
+    framework_name: str
+    def is_appropriate_type(self, tensor):
+        """helper method should recognize tensors it can handle"""
+        raise NotImplementedError()
+    def from_numpy(self, x):
+        raise NotImplementedError("framework doesn't support imperative execution")
+    def to_numpy(self, x):
+        raise NotImplementedError("framework doesn't support imperative execution")
+    def create_symbol(self, shape):
+        raise NotImplementedError("framework doesn't support symbolic computations")
+    def eval_symbol(self, symbol, symbol_value_pairs):
+        # symbol-value pairs is list[tuple[symbol, value-tensor]]
+        raise NotImplementedError("framework doesn't support symbolic computations")
+    def arange(self, start, stop):
+        # supplementary method used only in testing, so should implement CPU version
+        raise NotImplementedError("framework doesn't implement arange")
+    def shape(self, x):
+        """shape should return a tuple with integers or "shape symbols" (which will evaluate to actual size)"""
+        return x.shape
+    def reshape(self, x, shape):
+        return x.reshape(shape)
+    def transpose(self, x, axes):
+        return x.transpose(axes)
+    def reduce(self, x, operation, axes):
+        return getattr(x, operation)(axis=axes)
+    def stack_on_zeroth_dimension(self, tensors: list):
+        raise NotImplementedError()
+    def add_axis(self, x, new_position):
+        raise NotImplementedError()
+    def add_axes(self, x, n_axes, pos2len):
+        repeats = [1] * n_axes
+        for axis_position, axis_length in pos2len.items():
+            x = self.add_axis(x, axis_position)
+            repeats[axis_position] = axis_length
+        return self.tile(x, tuple(repeats))
+    def tile(self, x, repeats):
+        """repeats - same lengths as x.shape"""
+        raise NotImplementedError()
+    def concat(self, tensors, axis: int):
+        """concatenates tensors along axis.
+        Assume identical across tensors: devices, dtypes and shapes except selected axis."""
+        raise NotImplementedError()
+    def is_float_type(self, x):
+        # some backends (torch) can't compute average for non-floating types.
+        # Decided to drop average for all backends if type is not floating
+        raise NotImplementedError()
+    def layers(self):
+        raise NotImplementedError("backend does not provide layers")
+    def __repr__(self):
+        return "<einops backend for {}>".format(self.framework_name)
+    def einsum(self, pattern, *x):
+        raise NotImplementedError("backend does not support einsum")
+class UnknownSize:
+    """pseudo-symbol for symbolic frameworks which do not provide symbols for shape elements"""
+    def __floordiv__(self, other):
+        return self
+    def __eq__(self, other):
+        return True  # we don't know actual size
+    def __mul__(self, other):
+        return self
+    def __rmul__(self, other):
+        return self
+    def __hash__(self):
+        return hash(None)
+class NumpyBackend(AbstractBackend):
+    framework_name = "numpy"
+    def __init__(self):
+        import numpy
+        self.np = numpy
+    def is_appropriate_type(self, tensor):
+        return isinstance(tensor, self.np.ndarray)
+    def from_numpy(self, x):
+        return x
+    def to_numpy(self, x):
+        return x
+    def arange(self, start, stop):
+        return self.np.arange(start, stop)
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.np.stack(tensors)
+    def tile(self, x, repeats):
+        return self.np.tile(x, repeats)
+    def concat(self, tensors, axis: int):
+        return self.np.concatenate(tensors, axis=axis)
+    def is_float_type(self, x):
+        return x.dtype in ("float16", "float32", "float64", "float128", "bfloat16")
+    def add_axis(self, x, new_position):
+        return self.np.expand_dims(x, new_position)
+    def einsum(self, pattern, *x):
+        return self.np.einsum(pattern, *x)
+class JaxBackend(NumpyBackend):
+    framework_name = "jax"
+    def __init__(self):
+        super(JaxBackend, self).__init__()
+        self.onp = self.np
+        import jax.numpy
+        self.np = jax.numpy
+    def from_numpy(self, x):
+        return self.np.asarray(x)
+    def to_numpy(self, x):
+        return self.onp.asarray(x)
+class TorchBackend(AbstractBackend):
+    framework_name = "torch"
+    def __init__(self):
+        import torch
+        self.torch = torch
+        # importing would register operations in torch._dynamo for torch.compile
+        from . import _torch_specific  # noqa
+    def is_appropriate_type(self, tensor):
+        return isinstance(tensor, self.torch.Tensor)
+    def from_numpy(self, x):
+        variable = self.torch.from_numpy(x)
+        if self.is_float_type(variable):
+            # attach grad only to floating types
+            variable.requires_grad = True
+        return variable
+    def to_numpy(self, x):
+        return x.detach().cpu().numpy()
+    def arange(self, start, stop):
+        return self.torch.arange(start, stop, dtype=self.torch.int64)
+    def reduce(self, x, operation, reduced_axes):
+        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 in ("any", "all", "prod"):
+            # pytorch supports reducing only one operation at a time
+            for i in list(sorted(reduced_axes))[::-1]:
+                x = getattr(x, operation)(dim=i)
+            return x
+        else:
+            raise NotImplementedError("Unknown reduction ", operation)
+    def transpose(self, x, axes):
+        return x.permute(axes)
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.torch.stack(tensors)
+    def add_axes(self, x, n_axes, pos2len):
+        repeats = [-1] * n_axes
+        for axis_position, axis_length in pos2len.items():
+            x = self.add_axis(x, axis_position)
+            repeats[axis_position] = axis_length
+        return x.expand(repeats)
+    def tile(self, x, repeats):
+        return x.repeat(repeats)
+    def concat(self, tensors, axis: int):
+        return self.torch.cat(tensors, dim=axis)
+    def add_axis(self, x, new_position):
+        return self.torch.unsqueeze(x, new_position)
+    def is_float_type(self, x):
+        return x.dtype in [self.torch.float16, self.torch.float32, self.torch.float64, self.torch.bfloat16]
+    def layers(self):
+        from .layers import torch
+        return torch
+    def einsum(self, pattern, *x):
+        return self.torch.einsum(pattern, *x)
+class CupyBackend(AbstractBackend):
+    framework_name = "cupy"
+    def __init__(self):
+        import cupy
+        self.cupy = cupy
+    def is_appropriate_type(self, tensor):
+        return isinstance(tensor, self.cupy.ndarray)
+    def from_numpy(self, x):
+        return self.cupy.asarray(x)
+    def to_numpy(self, x):
+        return self.cupy.asnumpy(x)
+    def arange(self, start, stop):
+        return self.cupy.arange(start, stop)
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.cupy.stack(tensors)
+    def tile(self, x, repeats):
+        return self.cupy.tile(x, repeats)
+    def concat(self, tensors, axis: int):
+        return self.cupy.concatenate(tensors, axis=axis)
+    def add_axis(self, x, new_position):
+        return self.cupy.expand_dims(x, new_position)
+    def is_float_type(self, x):
+        return x.dtype in ("float16", "float32", "float64", "float128", "bfloat16")
+    def einsum(self, pattern, *x):
+        return self.cupy.einsum(pattern, *x)
+class HashableTuple:
+    """Overcomes non-hashability of symbolic elements"""
+    def __init__(self, elements: tuple):
+        self.elements = elements
+    def __iter__(self):
+        for x in self.elements:
+            yield x
+    def __len__(self):
+        return len(self.elements)
+    def __getitem__(self, item):
+        return self.elements[item]
+    # default equality and hash is used (True only with itself, hash taken of id)
+class TensorflowBackend(AbstractBackend):
+    framework_name = "tensorflow"
+    def __init__(self):
+        import tensorflow
+        self.tf = tensorflow
+    def is_appropriate_type(self, tensor):
+        return isinstance(tensor, (self.tf.Tensor, self.tf.Variable))
+    def from_numpy(self, x):
+        assert self.tf.executing_eagerly()
+        return self.tf.convert_to_tensor(x)
+    def to_numpy(self, x):
+        assert self.tf.executing_eagerly()
+        return x.numpy()
+    def arange(self, start, stop):
+        return self.tf.range(start, stop)
+    def shape(self, x):
+        if self.tf.executing_eagerly():
+            return tuple(UnknownSize() if d is None else int(d) for d in x.shape)
+        else:
+            static_shape = x.shape.as_list()
+            tf_shape = self.tf.shape(x)
+            # use the static shape where known, otherwise use the TF shape components
+            shape = tuple([s or tf_shape[dim] for dim, s in enumerate(static_shape)])
+            try:
+                hash(shape)
+                return shape
+            except BaseException:
+                # unhashable symbols in shape. Wrap tuple to be hashable.
+                return HashableTuple(shape)
+    def reduce(self, x, operation, axes):
+        return getattr(self.tf, "reduce_" + operation)(x, axis=axes)
+    def reshape(self, x, shape):
+        return self.tf.reshape(x, shape)
+    def transpose(self, x, axes):
+        return self.tf.transpose(x, axes)
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.tf.stack(tensors)
+    def tile(self, x, repeats):
+        return self.tf.tile(x, repeats)
+    def concat(self, tensors, axis: int):
+        return self.tf.concat(tensors, axis=axis)
+    def add_axis(self, x, new_position):
+        return self.tf.expand_dims(x, new_position)
+    def is_float_type(self, x):
+        return x.dtype in ("float16", "float32", "float64", "float128", "bfloat16")
+    def layers(self):
+        from .layers import tensorflow
+        return tensorflow
+    def einsum(self, pattern, *x):
+        return self.tf.einsum(pattern, *x)
+class TFKerasBackend(AbstractBackend):
+    framework_name = "tensorflow.keras"
+    def __init__(self):
+        import tensorflow as tf
+        self.tf = tf
+        self.keras = tf.keras
+        self.K = tf.keras.backend
+    def is_appropriate_type(self, tensor):
+        return self.tf.is_tensor(tensor) and self.K.is_keras_tensor(tensor)
+    def create_symbol(self, shape):
+        return self.keras.Input(batch_shape=shape)
+    def eval_symbol(self, symbol, symbol_value_pairs):
+        model = self.keras.models.Model([var for (var, _) in symbol_value_pairs], symbol)
+        return model.predict_on_batch([val for (_, val) in symbol_value_pairs])
+    def arange(self, start, stop):
+        return self.K.arange(start, stop)
+    def shape(self, x):
+        shape = self.K.shape(x)  # tf tensor
+        return HashableTuple(tuple(shape))
+    def reduce(self, x, operation, axes):
+        return getattr(self.K, operation)(x, axis=axes)
+    def reshape(self, x, shape):
+        return self.K.reshape(x, shape)
+    def transpose(self, x, axes):
+        return self.K.permute_dimensions(x, axes)
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.K.stack(tensors)
+    def tile(self, x, repeats):
+        return self.K.tile(x, repeats)
+    def concat(self, tensors, axis: int):
+        return self.K.concatenate(tensors, axis=axis)
+    def add_axis(self, x, new_position):
+        return self.K.expand_dims(x, new_position)
+    def is_float_type(self, x):
+        return "float" in self.K.dtype(x)
+    def layers(self):
+        from .layers import keras
+        return keras
+class OneFlowBackend(AbstractBackend):
+    framework_name = "oneflow"
+    def __init__(self):
+        import oneflow as flow
+        self.flow = flow
+    def is_appropriate_type(self, tensor):
+        return isinstance(tensor, self.flow.Tensor)
+    def from_numpy(self, x):
+        variable = self.flow.from_numpy(x)
+        if self.is_float_type(variable):
+            # attach grad only to floating types
+            variable.requires_grad = True
+        return variable
+    def to_numpy(self, x):
+        return x.detach().cpu().numpy()
+    def arange(self, start, stop):
+        return self.flow.arange(start, stop, dtype=self.flow.int64)
+    def reduce(self, x, operation, reduced_axes):
+        for axis in sorted(reduced_axes, reverse=True):
+            if operation == "min":
+                x, _ = x.min(dim=axis)
+            elif operation == "max":
+                x, _ = x.max(dim=axis)
+            elif operation in ["sum", "mean", "prod", "any", "all"]:
+                x = getattr(x, operation)(dim=axis)
+            else:
+                raise NotImplementedError("Unknown reduction ", operation)
+        return x
+    def transpose(self, x, axes):
+        return x.permute(axes)
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.flow.stack(tensors)
+    def add_axes(self, x, n_axes, pos2len):
+        repeats = [-1] * n_axes
+        for axis_position, axis_length in pos2len.items():
+            x = self.add_axis(x, axis_position)
+            repeats[axis_position] = axis_length
+        return x.expand(*repeats)
+    def tile(self, x, repeats):
+        return x.repeat(repeats)
+    def concat(self, tensors, axis: int):
+        return self.flow.concat(tensors, dim=axis)
+    def add_axis(self, x, new_position):
+        return self.flow.unsqueeze(x, new_position)
+    def is_float_type(self, x):
+        return x.dtype in [self.flow.float16, self.flow.float32, self.flow.float64]
+    def layers(self):
+        from .layers import oneflow
+        return oneflow
+    def einsum(self, pattern, *x):
+        return self.flow.einsum(pattern, *x)
+class PaddleBackend(AbstractBackend):
+    framework_name = "paddle"
+    def __init__(self):
+        import paddle
+        self.paddle = paddle
+    def is_appropriate_type(self, tensor):
+        return self.paddle.is_tensor(tensor)
+    def from_numpy(self, x):
+        tensor = self.paddle.to_tensor(x)
+        tensor.stop_gradient = False
+        return tensor
+    def to_numpy(self, x):
+        return x.detach().numpy()
+    def arange(self, start, stop):
+        return self.paddle.arange(start, stop, dtype=self.paddle.int64)
+    def reduce(self, x, operation, axes):
+        if len(axes) == x.ndim:
+            # currently paddle returns 1d tensor instead of 0d
+            return super().reduce(x, operation, axes).squeeze(0)
+        else:
+            return super().reduce(x, operation, axes)
+    def transpose(self, x, axes):
+        return x.transpose(axes)
+    def add_axes(self, x, n_axes, pos2len):
+        repeats = [-1] * n_axes
+        for axis_position, axis_length in pos2len.items():
+            x = self.add_axis(x, axis_position)
+            repeats[axis_position] = axis_length
+        return x.expand(repeats)
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.paddle.stack(tensors)
+    def reshape(self, x, shape):
+        return x.reshape(shape)
+    def tile(self, x, repeats):
+        return x.tile(repeats)
+    def concat(self, tensors, axis: int):
+        return self.paddle.concat(tensors, axis=axis)
+    def add_axis(self, x, new_position):
+        return x.unsqueeze(new_position)
+    def is_float_type(self, x):
+        return x.dtype in [self.paddle.float16, self.paddle.float32, self.paddle.float64]
+    def layers(self):
+        from .layers import paddle
+        return paddle
+    def einsum(self, pattern, *x):
+        return self.paddle.einsum(pattern, *x)
+    def shape(self, x):
+        return tuple(x.shape)
+class TinygradBackend(AbstractBackend):
+    framework_name = "tinygrad"
+    def __init__(self):
+        import tinygrad
+        self.tinygrad = tinygrad
+    def is_appropriate_type(self, tensor):
+        return isinstance(tensor, self.tinygrad.Tensor)
+    def from_numpy(self, x):
+        return self.tinygrad.Tensor(x)
+    def to_numpy(self, x):
+        return x.numpy()
+    def arange(self, start, stop):
+        return self.tinygrad.Tensor.arange(start, stop)
+    def shape(self, x):
+        return x.shape
+    def reshape(self, x, shape):
+        return x.reshape(shape)
+    def transpose(self, x, axes):
+        return x.permute(axes)
+    def reduce(self, x, operation, axes):
+        for axis in sorted(axes, reverse=True):
+            x = getattr(x, operation)(axis=axis)
+        return x
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.tinygrad.Tensor.stack(tensors)
+    def add_axis(self, x, new_position):
+        return x.unsqueeze(new_position)
+    def tile(self, x, repeats):
+        return x.repeat(repeats)
+    def concat(self, tensors, axis: int):
+        return tensors[0].cat(*tensors[1:], dim=axis) if len(tensors) > 1 else tensors[0]
+    def is_float_type(self, x):
+        return self.tinygrad.dtypes.is_float(x.dtype)
+    def einsum(self, pattern, *x):
+        return self.tinygrad.Tensor.einsum(pattern, *x)
+class PyTensorBackend(AbstractBackend):
+    framework_name = "pytensor"
+    def __init__(self):
+        from pytensor import tensor
+        self.pt = tensor
+    def is_appropriate_type(self, tensor):
+        return isinstance(tensor, self.pt.TensorVariable)
+    def is_float_type(self, x):
+        return x.dtype in self.pt.type.float_dtypes
+    def from_numpy(self, x):
+        return self.pt.as_tensor(x)
+    def to_numpy(self, x):
+        return x.eval()  # Will only work if there are no symbolic inputs
+    def create_symbol(self, shape):
+        if not isinstance(shape, tuple | list):
+            shape = (shape,)
+        return self.pt.tensor(shape=shape)
+    def eval_symbol(self, symbol, symbol_value_pairs):
+        return symbol.eval(dict(symbol_value_pairs))
+    def arange(self, start, stop):
+        return self.pt.arange(start, stop)
+    def shape(self, x):
+        # use the static shape dimensions where known
+        return tuple(
+            static_dim if static_dim is not None else symbolic_dim
+            for static_dim, symbolic_dim in zip(x.type.shape, x.shape)
+        )
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.pt.stack(tensors)
+    def tile(self, x, repeats):
+        return self.pt.tile(x, repeats)
+    def concat(self, tensors, axis: int):
+        return self.pt.concatenate(tensors, axis=axis)
+    def add_axis(self, x, new_position):
+        return self.pt.expand_dims(x, new_position)
+    def einsum(self, pattern, *x):
+        return self.pt.einsum(pattern, *x)

build/torch-cuda/_ops.py ADDED Viewed

	@@ -0,0 +1,38 @@

+import torch
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+    if hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+def _find_ops_name() -> str:
+    kernel_name = "einops"
+    unique_id = "d45adda"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+_OPS_NAME = _find_ops_name()
+ops = getattr(torch.ops, _OPS_NAME)
+def add_op_namespace_prefix(op_name: str) -> str:
+    """
+    Prefix op by namespace.
+    """
+    return f"{_OPS_NAME}::{op_name}"

build/torch-cuda/_torch_specific.py ADDED Viewed

	@@ -0,0 +1,128 @@

+"""
+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 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)
+# mirrors einops.einops._apply_recipe
+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":
+        # torch._dynamo and torch.compile appear in pytorch 2.0
+        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)
+    # CF: https://github.com/pytorch/pytorch/blob/2df939aacac68e9621fbd5d876c78d86e72b41e2/torch/_dynamo/__init__.py#L222
+    global _ops_were_registered_in_torchdynamo
+    _ops_were_registered_in_torchdynamo = True
+# module import automatically registers ops in torchdynamo
+allow_ops_in_compiled_graph()

build/torch-cuda/array_api.py ADDED Viewed

	@@ -0,0 +1,124 @@

+from typing import List, Tuple, Sequence
+from .einops import Tensor, Reduction, EinopsError, _prepare_transformation_recipe, _apply_recipe_array_api
+from .packing import analyze_pattern, prod
+def reduce(tensor: Tensor, pattern: str, reduction: Reduction, **axes_lengths: int) -> Tensor:
+    if isinstance(tensor, list):
+        if len(tensor) == 0:
+            raise TypeError("Einops can't be applied to an empty list")
+        xp = tensor[0].__array_namespace__()
+        tensor = xp.stack(tensor)
+    else:
+        xp = tensor.__array_namespace__()
+    try:
+        hashable_axes_lengths = tuple(axes_lengths.items())
+        recipe = _prepare_transformation_recipe(pattern, reduction, axes_names=tuple(axes_lengths), ndim=tensor.ndim)
+        return _apply_recipe_array_api(
+            xp,
+            recipe=recipe,
+            tensor=tensor,
+            reduction_type=reduction,
+            axes_lengths=hashable_axes_lengths,
+        )
+    except EinopsError as e:
+        message = ' Error while processing {}-reduction pattern "{}".'.format(reduction, pattern)
+        if not isinstance(tensor, list):
+            message += "\n Input tensor shape: {}. ".format(tensor.shape)
+        else:
+            message += "\n Input is list. "
+        message += "Additional info: {}.".format(axes_lengths)
+        raise EinopsError(message + "\n {}".format(e))
+def repeat(tensor: Tensor, pattern: str, **axes_lengths) -> Tensor:
+    return reduce(tensor, pattern, reduction="repeat", **axes_lengths)
+def rearrange(tensor: Tensor, pattern: str, **axes_lengths) -> Tensor:
+    return reduce(tensor, pattern, reduction="rearrange", **axes_lengths)
+def asnumpy(tensor: Tensor):
+    import numpy as np
+    return np.from_dlpack(tensor)
+Shape = Tuple
+def pack(tensors: Sequence[Tensor], pattern: str) -> Tuple[Tensor, List[Shape]]:
+    n_axes_before, n_axes_after, min_axes = analyze_pattern(pattern, "pack")
+    xp = tensors[0].__array_namespace__()
+    reshaped_tensors: List[Tensor] = []
+    packed_shapes: List[Shape] = []
+    for i, tensor in enumerate(tensors):
+        shape = tensor.shape
+        if len(shape) < min_axes:
+            raise EinopsError(
+                f"packed tensor #{i} (enumeration starts with 0) has shape {shape}, "
+                f"while pattern {pattern} assumes at least {min_axes} axes"
+            )
+        axis_after_packed_axes = len(shape) - n_axes_after
+        packed_shapes.append(shape[n_axes_before:axis_after_packed_axes])
+        reshaped_tensors.append(xp.reshape(tensor, (*shape[:n_axes_before], -1, *shape[axis_after_packed_axes:])))
+    return xp.concat(reshaped_tensors, axis=n_axes_before), packed_shapes
+def unpack(tensor: Tensor, packed_shapes: List[Shape], pattern: str) -> List[Tensor]:
+    xp = tensor.__array_namespace__()
+    n_axes_before, n_axes_after, min_axes = analyze_pattern(pattern, opname="unpack")
+    # backend = get_backend(tensor)
+    input_shape = tensor.shape
+    if len(input_shape) != n_axes_before + 1 + n_axes_after:
+        raise EinopsError(f"unpack(..., {pattern}) received input of wrong dim with shape {input_shape}")
+    unpacked_axis: int = n_axes_before
+    lengths_of_composed_axes: List[int] = [-1 if -1 in p_shape else prod(p_shape) for p_shape in packed_shapes]
+    n_unknown_composed_axes = sum(x == -1 for x in lengths_of_composed_axes)
+    if n_unknown_composed_axes > 1:
+        raise EinopsError(
+            f"unpack(..., {pattern}) received more than one -1 in {packed_shapes} and can't infer dimensions"
+        )
+    # following manipulations allow to skip some shape verifications
+    # and leave it to backends
+    # [[], [2, 3], [4], [-1, 5], [6]] < examples of packed_axis
+    # split positions when computed should be
+    # [0,   1,      7,   11,      N-6 , N ], where N = length of axis
+    split_positions = [0] * len(packed_shapes) + [input_shape[unpacked_axis]]
+    if n_unknown_composed_axes == 0:
+        for i, x in enumerate(lengths_of_composed_axes[:-1]):
+            split_positions[i + 1] = split_positions[i] + x
+    else:
+        unknown_composed_axis: int = lengths_of_composed_axes.index(-1)
+        for i in range(unknown_composed_axis):
+            split_positions[i + 1] = split_positions[i] + lengths_of_composed_axes[i]
+        for j in range(unknown_composed_axis + 1, len(lengths_of_composed_axes))[::-1]:
+            split_positions[j] = split_positions[j + 1] - lengths_of_composed_axes[j]
+    shape_start = input_shape[:unpacked_axis]
+    shape_end = input_shape[unpacked_axis + 1 :]
+    slice_filler = (slice(None, None),) * unpacked_axis
+    try:
+        return [
+            xp.reshape(
+                # shortest way slice arbitrary axis
+                tensor[(*slice_filler, slice(split_positions[i], split_positions[i + 1]), ...)],
+                (*shape_start, *element_shape, *shape_end),
+            )
+            for i, element_shape in enumerate(packed_shapes)
+        ]
+    except Exception:
+        # this hits if there is an error during reshapes, which means passed shapes were incorrect
+        raise RuntimeError(
+            f'Error during unpack(..., "{pattern}"): could not split axis of size {split_positions[-1]}'
+            f" into requested {packed_shapes}"
+        )

build/torch-cuda/einops.py ADDED Viewed

	@@ -0,0 +1,916 @@

+import functools
+import itertools
+import string
+import typing
+from collections import OrderedDict
+from typing import Set, Tuple, List, Dict, Union, Callable, Optional, TypeVar, cast, Any
+if typing.TYPE_CHECKING:
+    # for docstrings in pycharm
+    import numpy as np  # noqa E401
+from . import EinopsError
+from ._backends import get_backend
+from .parsing import ParsedExpression, _ellipsis, AnonymousAxis
+Tensor = TypeVar("Tensor")
+ReductionCallable = Callable[[Tensor, Tuple[int, ...]], Tensor]
+Reduction = Union[str, ReductionCallable]
+Size = typing.Any
+_reductions = ("min", "max", "sum", "mean", "prod", "any", "all")
+# magic integers are required to stay within
+# traceable subset of language
+_unknown_axis_length = -999999
+_expected_axis_length = -99999
+def _product(sequence: List[int]) -> int:
+    """minimalistic product that works both with numbers and symbols. Supports empty lists"""
+    result = 1
+    for element in sequence:
+        result *= element
+    return result
+def _reduce_axes(tensor, reduction_type: Reduction, reduced_axes: List[int], backend):
+    if callable(reduction_type):
+        # custom callable
+        return reduction_type(tensor, tuple(reduced_axes))
+    else:
+        # one of built-in operations
+        assert reduction_type in _reductions
+        if reduction_type == "mean":
+            if not backend.is_float_type(tensor):
+                raise NotImplementedError("reduce_mean is not available for non-floating tensors")
+        return backend.reduce(tensor, reduction_type, tuple(reduced_axes))
+def _optimize_transformation(init_shapes, reduced_axes, axes_reordering, final_shapes):
+    # 'collapses' neighboring axes if those participate in the result pattern in the same order
+    # TODO add support for added_axes
+    assert len(axes_reordering) + len(reduced_axes) == len(init_shapes)
+    # joining consecutive axes that will be reduced
+    # possibly we can skip this if all backends can optimize this (not sure)
+    reduced_axes = tuple(sorted(reduced_axes))
+    for i in range(len(reduced_axes) - 1)[::-1]:
+        if reduced_axes[i] + 1 == reduced_axes[i + 1]:
+            removed_axis = reduced_axes[i + 1]
+            removed_length = init_shapes[removed_axis]
+            init_shapes = init_shapes[:removed_axis] + init_shapes[removed_axis + 1 :]
+            init_shapes[removed_axis - 1] *= removed_length
+            reduced_axes = reduced_axes[: i + 1] + tuple(axis - 1 for axis in reduced_axes[i + 2 :])
+    # removing axes that are moved together during reshape
+    def build_mapping():
+        init_to_final = {}
+        for axis in range(len(init_shapes)):
+            if axis in reduced_axes:
+                init_to_final[axis] = None
+            else:
+                after_reduction = sum(x is not None for x in init_to_final.values())
+                init_to_final[axis] = list(axes_reordering).index(after_reduction)
+        return init_to_final
+    init_axis_to_final_axis = build_mapping()
+    for init_axis in range(len(init_shapes) - 1)[::-1]:
+        if init_axis_to_final_axis[init_axis] is None:
+            continue
+        if init_axis_to_final_axis[init_axis + 1] is None:
+            continue
+        if init_axis_to_final_axis[init_axis] + 1 == init_axis_to_final_axis[init_axis + 1]:
+            removed_axis = init_axis + 1
+            removed_length = init_shapes[removed_axis]
+            removed_axis_after_reduction = sum(x not in reduced_axes for x in range(removed_axis))
+            reduced_axes = tuple(axis if axis < removed_axis else axis - 1 for axis in reduced_axes)
+            init_shapes = init_shapes[:removed_axis] + init_shapes[removed_axis + 1 :]
+            init_shapes[removed_axis - 1] *= removed_length
+            old_reordering = axes_reordering
+            axes_reordering = []
+            for axis in old_reordering:
+                if axis == removed_axis_after_reduction:
+                    pass
+                elif axis < removed_axis_after_reduction:
+                    axes_reordering.append(axis)
+                else:
+                    axes_reordering.append(axis - 1)
+            init_axis_to_final_axis = build_mapping()
+    return init_shapes, reduced_axes, axes_reordering, final_shapes
+CookedRecipe = Tuple[Optional[List[int]], Optional[List[int]], List[int], Dict[int, int], Optional[List[int]], int]
+# Actual type is tuple[tuple[str, int], ...]
+# However torch.jit.script does not "understand" the correct type,
+# and torch_specific will use list version.
+HashableAxesLengths = Tuple[Tuple[str, int], ...]
+FakeHashableAxesLengths = List[Tuple[str, int]]
+class TransformRecipe:
+    """
+    Recipe describes actual computation pathway.
+    Recipe can be applied to a tensor or variable.
+    """
+    # structure is non-mutable. In future, this can be non-mutable dataclass (python 3.7+)
+    # update: pytorch 2.0 torch.jit.script seems to have problems with dataclasses unless they were explicitly provided
+    def __init__(
+        self,
+        # list of sizes (or just sizes) for elementary axes as they appear in left expression.
+        # this is what (after computing unknown parts) will be a shape after first transposition.
+        # This does not include any ellipsis dimensions.
+        elementary_axes_lengths: List[int],
+        # if additional axes are provided, they should be set in prev array
+        # This shows mapping from name to position
+        axis_name2elementary_axis: Dict[str, int],
+        # each dimension in input can help to reconstruct length of one elementary axis
+        # or verify one of dimensions. Each element points to element of elementary_axes_lengths.
+        input_composition_known_unknown: List[Tuple[List[int], List[int]]],
+        # permutation applied to elementary axes, if ellipsis is absent
+        axes_permutation: List[int],
+        # permutation puts reduced axes in the end, we only need to know the first position.
+        first_reduced_axis: int,
+        # at which positions which of elementary axes should appear. Axis position -> axis index.
+        added_axes: Dict[int, int],
+        # ids of axes as they appear in result, again pointers to elementary_axes_lengths,
+        # only used to infer result dimensions
+        output_composite_axes: List[List[int]],
+    ):
+        self.elementary_axes_lengths: List[int] = elementary_axes_lengths
+        self.axis_name2elementary_axis: Dict[str, int] = axis_name2elementary_axis
+        self.input_composition_known_unknown: List[Tuple[List[int], List[int]]] = input_composition_known_unknown
+        self.axes_permutation: List[int] = axes_permutation
+        self.first_reduced_axis: int = first_reduced_axis
+        self.added_axes: Dict[int, int] = added_axes
+        self.output_composite_axes: List[List[int]] = output_composite_axes
+def _reconstruct_from_shape_uncached(
+    self: TransformRecipe, shape: List[int], axes_dims: FakeHashableAxesLengths
+) -> CookedRecipe:
+    """
+    Reconstruct all actual parameters using shape.
+    Shape is a tuple that may contain integers, shape symbols (tf, theano) and UnknownSize (tf, previously mxnet)
+    known axes can be integers or symbols, but not Nones.
+    """
+    # magic number
+    need_init_reshape = False
+    # last axis is allocated for collapsed ellipsis
+    axes_lengths: List[int] = list(self.elementary_axes_lengths)
+    for axis, dim in axes_dims:
+        axes_lengths[self.axis_name2elementary_axis[axis]] = dim
+    for input_axis, (known_axes, unknown_axes) in enumerate(self.input_composition_known_unknown):
+        length = shape[input_axis]
+        if len(known_axes) == 0 and len(unknown_axes) == 1:
+            # shortcut for the most common case
+            axes_lengths[unknown_axes[0]] = length
+            continue
+        known_product = 1
+        for axis in known_axes:
+            known_product *= axes_lengths[axis]
+        if len(unknown_axes) == 0:
+            if isinstance(length, int) and isinstance(known_product, int) and length != known_product:
+                raise EinopsError(f"Shape mismatch, {length} != {known_product}")
+        else:
+            # assert len(unknown_axes) == 1, 'this is enforced when recipe is created, so commented out'
+            if isinstance(length, int) and isinstance(known_product, int) and length % known_product != 0:
+                raise EinopsError(f"Shape mismatch, can't divide axis of length {length} in chunks of {known_product}")
+            unknown_axis = unknown_axes[0]
+            inferred_length: int = length // known_product
+            axes_lengths[unknown_axis] = inferred_length
+        if len(known_axes) + len(unknown_axes) != 1:
+            need_init_reshape = True
+    # at this point all axes_lengths are computed (either have values or variables, but not Nones)
+    # elementary axes are ordered as they appear in input, then all added axes
+    init_shapes: Optional[List[int]] = axes_lengths[: len(self.axes_permutation)] if need_init_reshape else None
+    need_final_reshape = False
+    final_shapes: List[int] = []
+    for grouping in self.output_composite_axes:
+        lengths = [axes_lengths[elementary_axis] for elementary_axis in grouping]
+        final_shapes.append(_product(lengths))
+        if len(lengths) != 1:
+            need_final_reshape = True
+    added_axes: Dict[int, int] = {
+        pos: axes_lengths[pos_in_elementary] for pos, pos_in_elementary in self.added_axes.items()
+    }
+    # this list can be empty
+    reduced_axes = list(range(self.first_reduced_axis, len(self.axes_permutation)))
+    n_axes_after_adding_axes = len(added_axes) + len(self.axes_permutation)
+    axes_reordering: Optional[List[int]] = self.axes_permutation
+    if self.axes_permutation == list(range(len(self.axes_permutation))):
+        axes_reordering = None
+    _final_shapes = final_shapes if need_final_reshape else None
+    return init_shapes, axes_reordering, reduced_axes, added_axes, _final_shapes, n_axes_after_adding_axes
+_reconstruct_from_shape = functools.lru_cache(1024)(_reconstruct_from_shape_uncached)
+def _apply_recipe(
+    backend, recipe: TransformRecipe, tensor: Tensor, reduction_type: Reduction, axes_lengths: HashableAxesLengths
+) -> Tensor:
+    # this method implements actual work for all backends for 3 operations
+    try:
+        init_shapes, axes_reordering, reduced_axes, added_axes, final_shapes, n_axes_w_added = _reconstruct_from_shape(
+            recipe, backend.shape(tensor), axes_lengths
+        )
+    except TypeError:
+        # shape or one of passed axes lengths is not hashable (i.e. they are symbols)
+        _result = _reconstruct_from_shape_uncached(recipe, backend.shape(tensor), axes_lengths)
+        (init_shapes, axes_reordering, reduced_axes, added_axes, final_shapes, n_axes_w_added) = _result
+    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 = _reduce_axes(tensor, reduction_type=reduction_type, reduced_axes=reduced_axes, backend=backend)
+    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 _apply_recipe_array_api(
+    xp, recipe: TransformRecipe, tensor: Tensor, reduction_type: Reduction, axes_lengths: HashableAxesLengths
+) -> Tensor:
+    # completely-inline implementation
+    init_shapes, axes_reordering, reduced_axes, added_axes, final_shapes, n_axes_w_added = _reconstruct_from_shape(
+        recipe, tensor.shape, axes_lengths
+    )
+    if init_shapes is not None:
+        tensor = xp.reshape(tensor, init_shapes)
+    if axes_reordering is not None:
+        tensor = xp.permute_dims(tensor, axes_reordering)
+    if len(reduced_axes) > 0:
+        if callable(reduction_type):
+            # custom callable
+            tensor = reduction_type(tensor, tuple(reduced_axes))
+        else:
+            # one of built-in operations
+            assert reduction_type in _reductions
+            tensor = getattr(xp, reduction_type)(tensor, axis=tuple(reduced_axes))
+    if len(added_axes) > 0:
+        # we use broadcasting
+        for axis_position, axis_length in added_axes.items():
+            tensor = xp.expand_dims(tensor, axis=axis_position)
+        final_shape = list(tensor.shape)
+        for axis_position, axis_length in added_axes.items():
+            final_shape[axis_position] = axis_length
+        tensor = xp.broadcast_to(tensor, final_shape)
+    if final_shapes is not None:
+        tensor = xp.reshape(tensor, final_shapes)
+    return tensor
+@functools.lru_cache(256)
+def _prepare_transformation_recipe(
+    pattern: str,
+    operation: Reduction,
+    axes_names: Tuple[str, ...],
+    ndim: int,
+) -> TransformRecipe:
+    """Perform initial parsing of pattern and provided supplementary info
+    axes_lengths is a tuple of tuples (axis_name, axis_length)
+    """
+    left_str, rght_str = pattern.split("->")
+    left = ParsedExpression(left_str)
+    rght = ParsedExpression(rght_str)
+    # checking that axes are in agreement - new axes appear only in repeat, while disappear only in reduction
+    if not left.has_ellipsis and rght.has_ellipsis:
+        raise EinopsError("Ellipsis found in right side, but not left side of a pattern {}".format(pattern))
+    if left.has_ellipsis and left.has_ellipsis_parenthesized:
+        raise EinopsError("Ellipsis inside parenthesis in the left side is not allowed: {}".format(pattern))
+    if operation == "rearrange":
+        if left.has_non_unitary_anonymous_axes or rght.has_non_unitary_anonymous_axes:
+            raise EinopsError("Non-unitary anonymous axes are not supported in rearrange (exception is length 1)")
+        difference = set.symmetric_difference(left.identifiers, rght.identifiers)
+        if len(difference) > 0:
+            raise EinopsError("Identifiers only on one side of expression (should be on both): {}".format(difference))
+    elif operation == "repeat":
+        difference = set.difference(left.identifiers, rght.identifiers)
+        if len(difference) > 0:
+            raise EinopsError("Unexpected identifiers on the left side of repeat: {}".format(difference))
+        axes_without_size = set.difference(
+            {ax for ax in rght.identifiers if not isinstance(ax, AnonymousAxis)},
+            {*left.identifiers, *axes_names},
+        )
+        if len(axes_without_size) > 0:
+            raise EinopsError("Specify sizes for new axes in repeat: {}".format(axes_without_size))
+    elif operation in _reductions or callable(operation):
+        difference = set.difference(rght.identifiers, left.identifiers)
+        if len(difference) > 0:
+            raise EinopsError("Unexpected identifiers on the right side of reduce {}: {}".format(operation, difference))
+    else:
+        raise EinopsError("Unknown reduction {}. Expect one of {}.".format(operation, _reductions))
+    if left.has_ellipsis:
+        n_other_dims = len(left.composition) - 1
+        if ndim < n_other_dims:
+            raise EinopsError(f"Wrong shape: expected >={n_other_dims} dims. Received {ndim}-dim tensor.")
+        ellipsis_ndim = ndim - n_other_dims
+        ell_axes = [_ellipsis + str(i) for i in range(ellipsis_ndim)]
+        left_composition = []
+        for composite_axis in left.composition:
+            if composite_axis == _ellipsis:
+                for axis in ell_axes:
+                    left_composition.append([axis])
+            else:
+                left_composition.append(composite_axis)
+        rght_composition = []
+        for composite_axis in rght.composition:
+            if composite_axis == _ellipsis:
+                for axis in ell_axes:
+                    rght_composition.append([axis])
+            else:
+                group = []
+                for axis in composite_axis:
+                    if axis == _ellipsis:
+                        group.extend(ell_axes)
+                    else:
+                        group.append(axis)
+                rght_composition.append(group)
+        left.identifiers.update(ell_axes)
+        left.identifiers.remove(_ellipsis)
+        if rght.has_ellipsis:
+            rght.identifiers.update(ell_axes)
+            rght.identifiers.remove(_ellipsis)
+    else:
+        if ndim != len(left.composition):
+            raise EinopsError(f"Wrong shape: expected {len(left.composition)} dims. Received {ndim}-dim tensor.")
+        left_composition = left.composition
+        rght_composition = rght.composition
+    # parsing all dimensions to find out lengths
+    axis_name2known_length: Dict[Union[str, AnonymousAxis], int] = OrderedDict()
+    for composite_axis in left_composition:
+        for axis_name in composite_axis:
+            if isinstance(axis_name, AnonymousAxis):
+                axis_name2known_length[axis_name] = axis_name.value
+            else:
+                axis_name2known_length[axis_name] = _unknown_axis_length
+    # axis_ids_after_first_reshape = range(len(axis_name2known_length)) at this point
+    repeat_axes_names = []
+    for axis_name in rght.identifiers:
+        if axis_name not in axis_name2known_length:
+            if isinstance(axis_name, AnonymousAxis):
+                axis_name2known_length[axis_name] = axis_name.value
+            else:
+                axis_name2known_length[axis_name] = _unknown_axis_length
+            repeat_axes_names.append(axis_name)
+    axis_name2position = {name: position for position, name in enumerate(axis_name2known_length)}
+    # axes provided as kwargs
+    for elementary_axis in axes_names:
+        if not ParsedExpression.check_axis_name(elementary_axis):
+            raise EinopsError("Invalid name for an axis", elementary_axis)
+        if elementary_axis not in axis_name2known_length:
+            raise EinopsError("Axis {} is not used in transform".format(elementary_axis))
+        axis_name2known_length[elementary_axis] = _expected_axis_length
+    input_axes_known_unknown = []
+    # some shapes are inferred later - all information is prepared for faster inference
+    for i, composite_axis in enumerate(left_composition):
+        known: Set[str] = {axis for axis in composite_axis if axis_name2known_length[axis] != _unknown_axis_length}
+        unknown: Set[str] = {axis for axis in composite_axis if axis_name2known_length[axis] == _unknown_axis_length}
+        if len(unknown) > 1:
+            raise EinopsError("Could not infer sizes for {}".format(unknown))
+        assert len(unknown) + len(known) == len(composite_axis)
+        input_axes_known_unknown.append(
+            ([axis_name2position[axis] for axis in known], [axis_name2position[axis] for axis in unknown])
+        )
+    axis_position_after_reduction: Dict[str, int] = {}
+    for axis_name in itertools.chain(*left_composition):
+        if axis_name in rght.identifiers:
+            axis_position_after_reduction[axis_name] = len(axis_position_after_reduction)
+    result_axes_grouping: List[List[int]] = [
+        [axis_name2position[axis] for axis in composite_axis] for i, composite_axis in enumerate(rght_composition)
+    ]
+    ordered_axis_left = list(itertools.chain(*left_composition))
+    ordered_axis_rght = list(itertools.chain(*rght_composition))
+    reduced_axes = [axis for axis in ordered_axis_left if axis not in rght.identifiers]
+    order_after_transposition = [axis for axis in ordered_axis_rght if axis in left.identifiers] + reduced_axes
+    axes_permutation = [ordered_axis_left.index(axis) for axis in order_after_transposition]
+    added_axes = {
+        i: axis_name2position[axis_name]
+        for i, axis_name in enumerate(ordered_axis_rght)
+        if axis_name not in left.identifiers
+    }
+    first_reduced_axis = len(order_after_transposition) - len(reduced_axes)
+    return TransformRecipe(
+        elementary_axes_lengths=list(axis_name2known_length.values()),
+        axis_name2elementary_axis={axis: axis_name2position[axis] for axis in axes_names},
+        input_composition_known_unknown=input_axes_known_unknown,
+        axes_permutation=axes_permutation,
+        first_reduced_axis=first_reduced_axis,
+        added_axes=added_axes,
+        output_composite_axes=result_axes_grouping,
+    )
+def _prepare_recipes_for_all_dims(
+    pattern: str, operation: Reduction, axes_names: Tuple[str, ...]
+) -> Dict[int, TransformRecipe]:
+    """
+    Internal function, used in layers.
+    Layer makes all recipe creation when it is initialized, thus to keep recipes simple we pre-compute for all dims
+    """
+    left_str, rght_str = pattern.split("->")
+    left = ParsedExpression(left_str)
+    dims = [len(left.composition)]
+    if left.has_ellipsis:
+        dims = [len(left.composition) - 1 + ellipsis_dims for ellipsis_dims in range(8)]
+    return {ndim: _prepare_transformation_recipe(pattern, operation, axes_names, ndim=ndim) for ndim in dims}
+def reduce(tensor: Union[Tensor, List[Tensor]], pattern: str, reduction: Reduction, **axes_lengths: Size) -> Tensor:
+    """
+    einops.reduce combines rearrangement and reduction using reader-friendly notation.
+    Some examples:
+    ```python
+    >>> x = np.random.randn(100, 32, 64)
+    # perform max-reduction on the first axis
+    # Axis t does not appear on RHS - thus we reduced over t
+    >>> y = reduce(x, 't b c -> b c', 'max')
+    # same as previous, but using verbose names for axes
+    >>> y = reduce(x, 'time batch channel -> batch channel', 'max')
+    # let's pretend now that x is a batch of images
+    # with 4 dims: batch=10, height=20, width=30, channel=40
+    >>> x = np.random.randn(10, 20, 30, 40)
+    # 2d max-pooling with kernel size = 2 * 2 for image processing
+    >>> y1 = reduce(x, 'b c (h1 h2) (w1 w2) -> b c h1 w1', 'max', h2=2, w2=2)
+    # same as previous, using anonymous axes,
+    # note: only reduced axes can be anonymous
+    >>> y1 = reduce(x, 'b c (h1 2) (w1 2) -> b c h1 w1', 'max')
+    # adaptive 2d max-pooling to 3 * 4 grid,
+    # each element is max of 10x10 tile in the original tensor.
+    >>> reduce(x, 'b c (h1 h2) (w1 w2) -> b c h1 w1', 'max', h1=3, w1=4).shape
+    (10, 20, 3, 4)
+    # Global average pooling
+    >>> reduce(x, 'b c h w -> b c', 'mean').shape
+    (10, 20)
+    # subtracting mean over batch for each channel;
+    # similar to x - np.mean(x, axis=(0, 2, 3), keepdims=True)
+    >>> y = x - reduce(x, 'b c h w -> 1 c 1 1', 'mean')
+    # Subtracting per-image mean for each channel
+    >>> y = x - reduce(x, 'b c h w -> b c 1 1', 'mean')
+    # same as previous, but using empty compositions
+    >>> y = x - reduce(x, 'b c h w -> b c () ()', 'mean')
+    ```
+    Parameters:
+        tensor: tensor: tensor of any supported library (e.g. numpy.ndarray, tensorflow, pytorch).
+            list of tensors is also accepted, those should be of the same type and shape
+        pattern: string, reduction pattern
+        reduction: one of available reductions ('min', 'max', 'sum', 'mean', 'prod', 'any', 'all').
+            Alternatively, a callable f(tensor, reduced_axes) -> tensor can be provided.
+            This allows using various reductions like: np.max, np.nanmean, tf.reduce_logsumexp, torch.var, etc.
+        axes_lengths: any additional specifications for dimensions
+    Returns:
+        tensor of the same type as input
+    """
+    try:
+        if isinstance(tensor, list):
+            if len(tensor) == 0:
+                raise TypeError("Rearrange/Reduce/Repeat can't be applied to an empty list")
+            backend = get_backend(tensor[0])
+            tensor = backend.stack_on_zeroth_dimension(tensor)
+        else:
+            backend = get_backend(tensor)
+        hashable_axes_lengths = tuple(axes_lengths.items())
+        shape = backend.shape(tensor)
+        recipe = _prepare_transformation_recipe(pattern, reduction, axes_names=tuple(axes_lengths), ndim=len(shape))
+        return _apply_recipe(
+            backend, recipe, cast(Tensor, tensor), reduction_type=reduction, axes_lengths=hashable_axes_lengths
+        )
+    except EinopsError as e:
+        message = ' Error while processing {}-reduction pattern "{}".'.format(reduction, pattern)
+        if not isinstance(tensor, list):
+            message += "\n Input tensor shape: {}. ".format(shape)
+        else:
+            message += "\n Input is list. "
+        message += "Additional info: {}.".format(axes_lengths)
+        raise EinopsError(message + "\n {}".format(e))
+def rearrange(tensor: Union[Tensor, List[Tensor]], pattern: str, **axes_lengths: Size) -> Tensor:
+    """
+    einops.rearrange is a reader-friendly smart element reordering for multidimensional tensors.
+    This operation includes functionality of transpose (axes permutation), reshape (view), squeeze, unsqueeze,
+    stack, concatenate and other operations.
+    Examples:
+    ```python
+    # suppose we have a set of 32 images in "h w c" format (height-width-channel)
+    >>> images = [np.random.randn(30, 40, 3) for _ in range(32)]
+    # stack along first (batch) axis, output is a single array
+    >>> rearrange(images, 'b h w c -> b h w c').shape
+    (32, 30, 40, 3)
+    # stacked and reordered axes to "b c h w" format
+    >>> rearrange(images, 'b h w c -> b c h w').shape
+    (32, 3, 30, 40)
+    # concatenate images along height (vertical axis), 960 = 32 * 30
+    >>> rearrange(images, 'b h w c -> (b h) w c').shape
+    (960, 40, 3)
+    # concatenated images along horizontal axis, 1280 = 32 * 40
+    >>> rearrange(images, 'b h w c -> h (b w) c').shape
+    (30, 1280, 3)
+    # flattened each image into a vector, 3600 = 30 * 40 * 3
+    >>> rearrange(images, 'b h w c -> b (c h w)').shape
+    (32, 3600)
+    # split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), 128 = 32 * 2 * 2
+    >>> rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', h1=2, w1=2).shape
+    (128, 15, 20, 3)
+    # space-to-depth operation
+    >>> rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, w1=2).shape
+    (32, 15, 20, 12)
+    ```
+    When composing axes, C-order enumeration used (consecutive elements have different last axis).
+    Find more examples in einops tutorial.
+    Parameters:
+        tensor: tensor of any supported library (e.g. numpy.ndarray, tensorflow, pytorch).
+                list of tensors is also accepted, those should be of the same type and shape
+        pattern: string, rearrangement pattern
+        axes_lengths: any additional specifications for dimensions
+    Returns:
+        tensor of the same type as input. If possible, a view to the original tensor is returned.
+    """
+    return reduce(tensor, pattern, reduction="rearrange", **axes_lengths)
+def repeat(tensor: Union[Tensor, List[Tensor]], pattern: str, **axes_lengths: Size) -> Tensor:
+    """
+    einops.repeat allows reordering elements and repeating them in arbitrary combinations.
+    This operation includes functionality of repeat, tile, and broadcast functions.
+    Examples for repeat operation:
+    ```python
+    # a grayscale image (of shape height x width)
+    >>> image = np.random.randn(30, 40)
+    # change it to RGB format by repeating in each channel
+    >>> repeat(image, 'h w -> h w c', c=3).shape
+    (30, 40, 3)
+    # repeat image 2 times along height (vertical axis)
+    >>> repeat(image, 'h w -> (repeat h) w', repeat=2).shape
+    (60, 40)
+    # repeat image 2 time along height and 3 times along width
+    >>> repeat(image, 'h w -> (h2 h) (w3 w)', h2=2, w3=3).shape
+    (60, 120)
+    # convert each pixel to a small square 2x2. Upsample image by 2x
+    >>> repeat(image, 'h w -> (h h2) (w w2)', h2=2, w2=2).shape
+    (60, 80)
+    # pixelate image first by downsampling by 2x, then upsampling
+    >>> downsampled = reduce(image, '(h h2) (w w2) -> h w', 'mean', h2=2, w2=2)
+    >>> repeat(downsampled, 'h w -> (h h2) (w w2)', h2=2, w2=2).shape
+    (30, 40)
+    ```
+    When composing axes, C-order enumeration used (consecutive elements have different last axis).
+    Find more examples in einops tutorial.
+    Parameters:
+        tensor: tensor of any supported library (e.g. numpy.ndarray, tensorflow, pytorch).
+            list of tensors is also accepted, those should be of the same type and shape
+        pattern: string, rearrangement pattern
+        axes_lengths: any additional specifications for dimensions
+    Returns:
+        Tensor of the same type as input. If possible, a view to the original tensor is returned.
+    """
+    return reduce(tensor, pattern, reduction="repeat", **axes_lengths)
+def parse_shape(x: Tensor, pattern: str) -> dict:
+    """
+    Parse a tensor shape to dictionary mapping axes names to their lengths.
+    ```python
+    # Use underscore to skip the dimension in parsing.
+    >>> x = np.zeros([2, 3, 5, 7])
+    >>> parse_shape(x, 'batch _ h w')
+    {'batch': 2, 'h': 5, 'w': 7}
+    # `parse_shape` output can be used to specify axes_lengths for other operations:
+    >>> y = np.zeros([700])
+    >>> rearrange(y, '(b c h w) -> b c h w', **parse_shape(x, 'b _ h w')).shape
+    (2, 10, 5, 7)
+    ```
+    For symbolic frameworks may return symbols, not integers.
+    Parameters:
+        x: tensor of any supported framework
+        pattern: str, space separated names for axes, underscore means skip axis
+    Returns:
+        dict, maps axes names to their lengths
+    """
+    exp = ParsedExpression(pattern, allow_underscore=True)
+    shape = get_backend(x).shape(x)
+    if exp.has_composed_axes():
+        raise RuntimeError(f"Can't parse shape with composite axes: {pattern} {shape}")
+    if len(shape) != len(exp.composition):
+        if exp.has_ellipsis:
+            if len(shape) < len(exp.composition) - 1:
+                raise RuntimeError(f"Can't parse shape with this number of dimensions: {pattern} {shape}")
+        else:
+            raise RuntimeError(f"Can't parse shape with different number of dimensions: {pattern} {shape}")
+    if exp.has_ellipsis:
+        ellipsis_idx = exp.composition.index(_ellipsis)
+        composition = (
+            exp.composition[:ellipsis_idx]
+            + ["_"] * (len(shape) - len(exp.composition) + 1)
+            + exp.composition[ellipsis_idx + 1 :]
+        )
+    else:
+        composition = exp.composition
+    result = {}
+    for axes, axis_length in zip(composition, shape):  # type: ignore
+        # axes either [], or [AnonymousAxis] or ['axis_name']
+        if len(axes) == 0:
+            if axis_length != 1:
+                raise RuntimeError(f"Length of axis is not 1: {pattern} {shape}")
+        else:
+            [axis] = axes
+            if isinstance(axis, str):
+                if axis != "_":
+                    result[axis] = axis_length
+            else:
+                if axis.value != axis_length:
+                    raise RuntimeError(f"Length of anonymous axis does not match: {pattern} {shape}")
+    return result
+# _enumerate_directions is not exposed in the public API
+def _enumerate_directions(x):
+    """
+    For an n-dimensional tensor, returns tensors to enumerate each axis.
+    ```python
+    x = np.zeros([2, 3, 4]) # or any other tensor
+    i, j, k = _enumerate_directions(x)
+    result = i + 2*j + 3*k
+    ```
+    `result[i, j, k] = i + 2j + 3k`, and also has the same shape as result
+    Works very similarly to numpy.ogrid (open indexing grid)
+    """
+    backend = get_backend(x)
+    shape = backend.shape(x)
+    result = []
+    for axis_id, axis_length in enumerate(shape):
+        shape = [1] * len(shape)
+        shape[axis_id] = axis_length
+        result.append(backend.reshape(backend.arange(0, axis_length), shape))
+    return result
+# to avoid importing numpy
+np_ndarray = Any
+def asnumpy(tensor: Tensor) -> np_ndarray:
+    """
+    Convert a tensor of an imperative framework (i.e. numpy/cupy/torch/jax/etc.) to `numpy.ndarray`
+    Parameters:
+        tensor: tensor of any known imperative framework
+    Returns:
+        `numpy.ndarray`, converted to numpy
+    """
+    return get_backend(tensor).to_numpy(tensor)
+def _validate_einsum_axis_name(axis_name):
+    if len(axis_name) == 0:
+        raise NotImplementedError("Singleton () axes are not yet supported in einsum.")
+    if len(axis_name) > 1:
+        raise NotImplementedError("Shape rearrangement is not yet supported in einsum.")
+    axis_name = axis_name[0]
+    if isinstance(axis_name, AnonymousAxis):
+        raise NotImplementedError("Anonymous axes are not yet supported in einsum.")
+    if len(axis_name) == 0:
+        raise RuntimeError("Encountered empty axis name in einsum.")
+    if not isinstance(axis_name, str):
+        raise RuntimeError("Axis name in einsum must be a string.")
+@functools.lru_cache(256)
+def _compactify_pattern_for_einsum(pattern: str) -> str:
+    if "->" not in pattern:
+        # numpy allows this, so make sure users
+        # don't accidentally do something like this.
+        raise ValueError("Einsum pattern must contain '->'.")
+    lefts_str, right_str = pattern.split("->")
+    lefts = [ParsedExpression(left, allow_underscore=True, allow_duplicates=True) for left in lefts_str.split(",")]
+    right = ParsedExpression(right_str, allow_underscore=True)
+    # Start from 'a' and go up to 'Z'
+    output_axis_names = string.ascii_letters
+    i = 0
+    axis_name_mapping = {}
+    left_patterns = []
+    for left in lefts:
+        left_pattern = ""
+        for raw_axis_name in left.composition:
+            if raw_axis_name == _ellipsis:
+                left_pattern += "..."
+                continue
+            _validate_einsum_axis_name(raw_axis_name)
+            axis_name = raw_axis_name[0]
+            if axis_name not in axis_name_mapping:
+                if i >= len(output_axis_names):
+                    raise RuntimeError("Too many axes in einsum.")
+                axis_name_mapping[axis_name] = output_axis_names[i]
+                i += 1
+            left_pattern += axis_name_mapping[axis_name]
+        left_patterns.append(left_pattern)
+    compact_pattern = ",".join(left_patterns) + "->"
+    for raw_axis_name in right.composition:
+        if raw_axis_name == _ellipsis:
+            compact_pattern += "..."
+            continue
+        _validate_einsum_axis_name(raw_axis_name)
+        axis_name = raw_axis_name[0]
+        if axis_name not in axis_name_mapping:
+            raise EinopsError(f"Unknown axis {axis_name} on right side of einsum {pattern}.")
+        compact_pattern += axis_name_mapping[axis_name]
+    return compact_pattern
+@typing.overload
+def einsum(tensor: Tensor, pattern: str, /) -> Tensor: ...
+@typing.overload
+def einsum(tensor1: Tensor, tensor2: Tensor, pattern: str, /) -> Tensor: ...
+@typing.overload
+def einsum(tensor1: Tensor, tensor2: Tensor, tensor3: Tensor, pattern: str, /) -> Tensor: ...
+@typing.overload
+def einsum(tensor1: Tensor, tensor2: Tensor, tensor3: Tensor, tensor4: Tensor, pattern: str, /) -> Tensor: ...
+def einsum(*tensors_and_pattern: Union[Tensor, str]) -> Tensor:
+    r"""
+    einops.einsum calls einsum operations with einops-style named
+    axes indexing, computing tensor products with an arbitrary
+    number of tensors. Unlike typical einsum syntax, here you must
+    pass tensors first, and then the pattern.
+    Also, note that rearrange operations such as `"(batch chan) out"`,
+    or singleton axes `()`, are not currently supported.
+    Examples:
+    For a given pattern such as:
+    ```python
+    >>> x, y, z = np.random.randn(3, 20, 20, 20)
+    >>> output = einsum(x, y, z, "a b c, c b d, a g k -> a b k")
+    ```
+    the following formula is computed:
+    ```tex
+    output[a, b, k] =
+        \sum_{c, d, g} x[a, b, c] * y[c, b, d] * z[a, g, k]
+    ```
+    where the summation over `c`, `d`, and `g` is performed
+    because those axes names do not appear on the right-hand side.
+    Let's see some additional examples:
+    ```python
+    # Filter a set of images:
+    >>> batched_images = np.random.randn(128, 16, 16)
+    >>> filters = np.random.randn(16, 16, 30)
+    >>> result = einsum(batched_images, filters,
+    ...                 "batch h w, h w channel -> batch channel")
+    >>> result.shape
+    (128, 30)
+    # Matrix multiplication, with an unknown input shape:
+    >>> batch_shape = (50, 30)
+    >>> data = np.random.randn(*batch_shape, 20)
+    >>> weights = np.random.randn(10, 20)
+    >>> result = einsum(weights, data,
+    ...                 "out_dim in_dim, ... in_dim -> ... out_dim")
+    >>> result.shape
+    (50, 30, 10)
+    # Matrix trace on a single tensor:
+    >>> matrix = np.random.randn(10, 10)
+    >>> result = einsum(matrix, "i i ->")
+    >>> result.shape
+    ()
+    ```
+    Parameters:
+        tensors_and_pattern:
+            tensors: tensors of any supported library (numpy, tensorflow, pytorch, jax).
+            pattern: string, einsum pattern, with commas
+                separating specifications for each tensor.
+                pattern should be provided after all tensors.
+    Returns:
+        Tensor of the same type as input, after processing with einsum.
+    """
+    if len(tensors_and_pattern) <= 1:
+        raise ValueError(
+            "`einops.einsum` takes at minimum two arguments: the tensors (at least one), followed by the pattern."
+        )
+    pattern = tensors_and_pattern[-1]
+    if not isinstance(pattern, str):
+        raise ValueError(
+            "The last argument passed to `einops.einsum` must be a string, representing the einsum pattern."
+        )
+    tensors = tensors_and_pattern[:-1]
+    pattern = _compactify_pattern_for_einsum(pattern)
+    return get_backend(tensors[0]).einsum(pattern, *tensors)

build/torch-cuda/experimental/__init__.py ADDED Viewed

File without changes

build/torch-cuda/experimental/indexing.py ADDED Viewed

	@@ -0,0 +1,5 @@

+"""
+This file contained some thoughts on indexing.
+These ideas were developed further in eindex (separate package).
+"""

build/torch-cuda/layers/__init__.py ADDED Viewed

	@@ -0,0 +1,106 @@

+__author__ = "Alex Rogozhnikov"
+from typing import Any, Dict
+from ..einops import TransformRecipe, _apply_recipe, _prepare_recipes_for_all_dims, get_backend
+from .. import EinopsError
+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 += ", {}={}".format(axis, length)
+        return "{}({})".format(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(" Error while preparing {!r}\n {}".format(self, e))
+    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 = "{!r}, {!r}".format(self.pattern, self.reduction)
+        for axis, length in self.axes_lengths.items():
+            params += ", {}={}".format(axis, length)
+        return "{}({})".format(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(" Error while preparing {!r}\n {}".format(self, e))
+    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"])

build/torch-cuda/layers/_einmix.py ADDED Viewed

	@@ -0,0 +1,229 @@

+from typing import Any, List, Optional, Dict
+from .. import EinopsError
+from ..parsing import ParsedExpression, _ellipsis
+import warnings
+import string
+from ..einops import _product
+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("Dimension {} of weight should be specified".format(axis))
+        _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)")
+        _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 = list(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 = "{},{}->{}".format(
+            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 += ", {}={}".format(axis, length)
+        return "{}({})".format(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

build/torch-cuda/layers/flax.py ADDED Viewed

	@@ -0,0 +1,82 @@

+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

build/torch-cuda/layers/keras.py ADDED Viewed

	@@ -0,0 +1,9 @@

+__author__ = "Alex Rogozhnikov"
+from ..layers.tensorflow import Rearrange, Reduce, EinMix
+keras_custom_objects = {
+    Rearrange.__name__: Rearrange,
+    Reduce.__name__: Reduce,
+    EinMix.__name__: EinMix,
+}

build/torch-cuda/layers/oneflow.py ADDED Viewed

	@@ -0,0 +1,54 @@

+from typing import Optional, Dict, 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

build/torch-cuda/layers/paddle.py ADDED Viewed

	@@ -0,0 +1,58 @@

+from typing import Optional, Dict, 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

build/torch-cuda/layers/tensorflow.py ADDED Viewed

	@@ -0,0 +1,103 @@

+"""
+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 Optional, Dict, 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,
+        }

build/torch-cuda/layers/torch.py ADDED Viewed

	@@ -0,0 +1,67 @@

+from typing import Optional, Dict, cast
+import torch
+from . import RearrangeMixin, ReduceMixin
+from ._einmix import _EinmixMixin
+from .._torch_specific import apply_for_scriptable_torch
+__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

build/torch-cuda/metadata.json ADDED Viewed

	@@ -0,0 +1,45 @@

+{
+  "name": "einops",
+  "id": "_einops_cuda_d45adda",
+  "version": 1,
+  "license": "MIT",
+  "upstream": "https://github.com/arogozhnikov/einops.git",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Zi5UF5cdv0EjT90oh0zsv0Jfi4QbTrfZnE0Tv5on8F8=",
+      "_backends.py": "7BAiKWdBLLvAiryfhwDbwxglGDW9Dh5FNMT1ERSuzxk=",
+      "_ops.py": "11R2mMPUK6wLQIwJo8geCYiUgC/WM4OCjMiT15zK3hQ=",
+      "_torch_specific.py": "CxpGcgqlBn4BH+Qkpmr3dFBPTid1Rz/R0f9TGb36Bsw=",
+      "array_api.py": "jOb8RhwLS9wob/Y/e/KrnBR6ihQPoB2Ly0tfrHr+/Zk=",
+      "einops.py": "sXvD8SWFqufziyQJKRPmfAGHVN1cMDvYOPNuZ8L1XQU=",
+      "experimental/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=",
+      "experimental/indexing.py": "yFFflW3+kV6/5PPJU7/jOJsJBCWCWlE4dGlu9gwSPXo=",
+      "layers/__init__.py": "vBtnAt2afs4QlqpeFU4dlZNxBuC9IXl3fmilk+2OzHM=",
+      "layers/_einmix.py": "9cDMcCmn2y1jN9xLx587GAY1GCb+9TvEjJWaNNX4Vps=",
+      "layers/flax.py": "zFy83gSLRm31cLuKFRvZ82/HsefnXPbRvkKZh1KkC1I=",
+      "layers/keras.py": "+7So0w94phvf9HdW0xi2mSeBg02qVPvAyfp/1XR02NM=",
+      "layers/oneflow.py": "YEPzz4xc7BDRQfb8ulD3teqQJdbO6qQg7Z4KIPVTLz8=",
+      "layers/paddle.py": "8cRZQ8BT9vYEczh7pNProuTM/3XjLty2ht2sdvXNFiI=",
+      "layers/tensorflow.py": "T9uhSVwbXREahc31ARAHoN5K+7zsuS8NRNPdY6Zk1Bc=",
+      "layers/torch.py": "504G99kEgy7dk1UPBbj9hzJmZkAHwVhMDFN/8J+p3C8=",
+      "packing.py": "vBjwbVWs3OwmI83BMNoeu3jBFAbAiy10i3ClmZxZtCQ=",
+      "parsing.py": "tXcSr4W1mbePUu+oIYgG0cKYNxsPXuVYFIHynj1FCg4="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "d0610aa58db33b142c86b59598a2a1c730f52996",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "d45addadb0c380f3d7cc2813310b3bcd75f8aa9a",
+      "dirty": false
+    }
+  }
+}

build/torch-cuda/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUUacGOOXSbwomtiDehEauSQXdJwwwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODA3MTM0NDU2WhcNMjYwODA3MTM1NDU2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEWdnIdX2ky5B75yPQrw7yVzWMLKh4/+hxIVc7Xht7bvW9sWDfTFvAk9SI6ouqivPAi6lpo15e3xwvy8+vj2V5TqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUJ67AbLESuta0gOytiK2nhg62RT8wHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoZDQ1YWRkYWRiMGMzODBmM2Q3Y2MyODEzMzEwYjNiY2Q3NWY4YWE5YTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoZDQ1YWRkYWRiMGMzODBmM2Q3Y2MyODEzMzEwYjNiY2Q3NWY4YWE5YTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoZDQ1YWRkYWRiMGMzODBmM2Q3Y2MyODEzMzEwYjNiY2Q3NWY4YWE5YTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGQ0NWFkZGFkYjBjMzgwZjNkN2NjMjgxMzMxMGIzYmNkNzVmOGFhOWEwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzExODM4OTEwNDkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn9x4QXQAAAQDAEcwRQIhAN0G29OPXX93oRMXIydM/ZUDE1YY72DiGL7z+fNTV5LYAiAJeGPwHqGWVUGQdxIGkyHyALGPxESddGNwL4Lme8znYDAKBggqhkjOPQQDAwNoADBlAjBn62jhPKIHDtuErG2gU18qy5Ed+gyQ3E49hlgbvzWqw6LCpsdJujoG4o8Zm2PgJPkCMQDBbeoUjvA8MefZzoRWJYKlrdlc5cV3NEfxbPMW9KZHlXn4phUuGUOTvxdrgsq2YI4="},"tlogEntries":[{"logIndex":"2370478831","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1786110296","inclusionPromise":{"signedEntryTimestamp":"MEQCIA4ddvAnNNm4iLQdOUrp/YlF62yKHmcfUKouhmMp+53RAiB24MzS+UaRunfNspUoMSiKLE3CO+xPpNu7DVvA8XdROg=="},"inclusionProof":{"logIndex":"2248574569","rootHash":"xhF+Sk3KdHoNMadzxW/BHVnlSMPtvtFvzzedzMX7myo=","treeSize":"2248574572","hashes":["X2Iv4avwV+o67VMPuEKTbGrbWc2RxKOtpINPOr9TmaE=","Lc1oHGdjLyQYnP+8PeejD9qE8oi88gyyPzGwnwpZOGc=","jmc6FKLo0hnfILv8q3LOe6XkwB47NMfXrf+eE+NiC/s=","lG0vBfmy3NyKqsDg3rHUTYmwdKoqekKk/5yxaoIwaeE=","5IYm0izDiwos8kF2OODaxXQVp5GSAjhY7ysAOvm2xMs=","XXWXzYRrsBDaQeQZeD9JT3OtOuvg//LUnu+u3IwZFOQ=","1iszQpDQvYGSRDukIT2HFergOhdefiQIlIp9VIaKtco=","ByRm8lkUhtWxweoZD8P90jNd2or8eFiw04pJXBx1cHE=","BS5WEjUIQqZG8MjRFGLBhqD4ZyOzwjsu+LZ0qVkyKtU=","B+BvP8FCZx7ZE5eJJXnytVRt2HDakcIeugCRQ/HvxH8=","ILMoQSbIb83ZRi3LtJ9th99rqP+Za3UA534sa6mEYKg=","i5Zl8FZrDwxCDv2e2DNO2M8JvpR/c11ElvCZS53/teA=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2248574572\nxhF+Sk3KdHoNMadzxW/BHVnlSMPtvtFvzzedzMX7myo=\n\n— rekor.sigstore.dev wNI9ajBEAiBpZrbe/RIeUbWTuvS/ILM3yQ0DnjkfWhoVTpN9Wl7PzAIgTN+mqrnFaae6S6qkei8QpiYq/ftmleO9NU5qjYDjojs=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI2NDI3MzcwMmZmMDMwN2MwYjA2Y2E2MWMzOWE4ZjQ4YjE5OGFmM2RmNGE2ZGZiNjcwZGJiMDhkZjA4NjEwNjhkIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNEWUx0ZGhWWDNpMEdQeHNPMXM3VmlJNjFEbUM0cmNORFBoYnpNZXpqbnNBSWhBT0VJVWZJd3dGY1VkcytmM1N1R0UyYUZ1MTY3cGFXR3I5V2dreTg3Y2g2ayIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlZXRmpSMDlQV0ZOaWQyOXRkR2xFWldoRllYVlRVVmhrU25kM2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFUVROTlZFMHdUa1JWTWxkb1kwNU5hbGwzVDBSQk0wMVVUVEZPUkZVeVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZYWkc1SlpGZ3lhM2sxUWpjMWVWQlJjbmMzZVZaNlYwMU1TMmcwTHl0b2VFbFdZemNLV0doME4ySjJWemx6VjBSbVZFWjJRV3M1VTBrMmIzVnhhWFpRUVdrMmJIQnZNVFZsTTNoM2RuazRLM1pxTWxZMVZIRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZLTmpkQkNtSk1SVk4xZEdFd1owOTVkR2xMTW01b1p6WXlVbFE0ZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOWFSRkV4V1ZkU2ExbFhVbWxOUjAxNlQwUkNiVTB5VVROWk1rMTVDazlFUlhwTmVrVjNXV3BPYVZreVVUTk9WMWswV1ZkRk5WbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMXBFVVRGWlYxSnJXVmRTYVUxSFRYcFBSRUp0VFRKUk0xa3lUWGxQUkVWNlRYcEZkMWxxVG1rS1dUSlJNMDVYV1RSWlYwVTFXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMXBFVVRGWlYxSnJXVmRTYVUxSFRYcFBSRUp0VFRKUk0xa3lUWGtLVDBSRmVrMTZSWGRaYWs1cFdUSlJNMDVYV1RSWlYwVTFXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZFJNRTVYUm1zS1drZEdhMWxxUW1wTmVtZDNXbXBPYTA0eVRtcE5hbWQ0VFhwTmVFMUhTWHBaYlU1clRucFdiVTlIUm1oUFYwVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtWNFQwUk5ORTlVUlhkT1JHdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDVlRFJSV0ZGQlFVRlJSQXBCUldOM1VsRkphRUZPTUVjeU9VOVFXRmc1TTI5U1RWaEplV1JOTDFwVlJFVXhXVmszTWtScFIwdzNlaXRtVGxSV05VeFpRV2xCU21WSFVIZEljVWRYQ2xaVlIxRmtlRWxIYTNsSWVVRk1SMUI0UlZOa1pFZE9kMHcwVEcxbE9IcHVXVVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcENiall5YW1nS1VFdEpTRVIwZFVWeVJ6Sm5WVEU0Y1hrMVJXUXJaM2xSTTBVME9XaHNaMkoyZWxkeGR6Wk1RM0J6WkVwMWFtOUhORzg0V20weVVHZEtVR3REVFZGRVFncGlaVzlWYW5aQk9FMWxabHA2YjFKWFNsbExiSEprYkdNMVkxWXpUa1ZtZUdKUVRWYzVTMXBJYkZodU5IQm9WWFZIVlU5VWRuaGtjbWR6Y1RKWlNUUTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgbkhBcHaQNxUEmWxJK4CpVHFAmtNtaRXoaCGBUaQcr1kCFE6NZqPwuGlzKx+1Ps+/X3Ay94XGGA8yMDI2MDgwNzEzNDQ1NlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwODA3MTM0NDU2WjAvBgkqhkiG9w0BCQQxIgQg5Reu7QB5QfKf5T1fzU4O5t5c2c6k0fsObLh7Fx77v1MwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIvEPL1+OVrRjfwSbXEUergXdZI/72Vc6lh07pIjIG1W/uvhCuzJGAinnFUjeJG8HkCMQCUcOD5vmAgXv19OrIy3ZDqrqq324j70th+nIkTZFyWJgma6R6PHJtvVdtQrSa0xME="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"ZCc3Av8DB8CwbKYcOaj0ixmK899KbftnDbsI3whhBo0="},"signature":"MEYCIQCDYLtdhVX3i0GPxsO1s7ViI61DmC4rcNDPhbzMezjnsAIhAOEIUfIwwFcUds+f3SuGE2aFu167paWGr9Wgky87ch6k"}}

build/torch-cuda/packing.py ADDED Viewed

	@@ -0,0 +1,190 @@

+from functools import lru_cache
+from typing import List, Union, TypeVar, Tuple, Sequence
+from . import EinopsError
+from ._backends import get_backend
+from .parsing import ParsedExpression
+Tensor = TypeVar("Tensor")
+Shape = Union[Tuple[int, ...], List[int]]
+@lru_cache(maxsize=128)
+def analyze_pattern(pattern: str, opname: str) -> Tuple[int, int, int]:
+    # Maybe some validation of identifiers?
+    axes = pattern.split()
+    axes_set = set(axes)
+    if len(axes) != len(axes_set):
+        raise EinopsError(f'Duplicates in axes names in {opname}(..., "{pattern}")')
+    if "*" not in axes_set:
+        raise EinopsError(f'No *-axis in {opname}(..., "{pattern}")')
+    for axis in axes:
+        if axis != "*":
+            is_valid, reason = ParsedExpression.check_axis_name_return_reason(axis)
+            if not is_valid:
+                raise EinopsError(f'Invalid axis name {axis} in {opname}(..., "{pattern}")')
+    n_axes_before = axes.index("*")
+    n_axes_after = len(axes) - n_axes_before - 1
+    min_axes = n_axes_before + n_axes_after
+    return n_axes_before, n_axes_after, min_axes
+def pack(tensors: Sequence[Tensor], pattern: str) -> Tuple[Tensor, List[Shape]]:
+    """
+    Packs several tensors into one.
+    See einops tutorial for introduction into packing (and how it replaces stack and concatenation).
+    Parameters:
+        tensors: tensors to be packed, can be of different dimensionality
+        pattern: pattern that is shared for all inputs and output, e.g. "i j * k" or "batch seq *"
+    Returns:
+        (packed_tensor, packed_shapes aka PS)
+    Example:
+    ```python
+    >>> from numpy import zeros as Z
+    >>> inputs = [Z([2, 3, 5]), Z([2, 3, 7, 5]), Z([2, 3, 7, 9, 5])]
+    >>> packed, ps = pack(inputs, 'i j * k')
+    >>> packed.shape, ps
+    ((2, 3, 71, 5), [(), (7,), (7, 9)])
+    ```
+    In this example, axes were matched to: i=2, j=3, k=5 based on order (first, second, and last).
+    All other axes were 'packed' and concatenated.
+    PS (packed shapes) contains information about axes that were matched to '*' in every input.
+    Resulting tensor has as many elements as all inputs in total.
+    Packing can be reversed with unpack, which additionally needs PS (packed shapes) to reconstruct order.
+    ```python
+    >>> inputs_unpacked = unpack(packed, ps, 'i j * k')
+    >>> [x.shape for x in inputs_unpacked]
+    [(2, 3, 5), (2, 3, 7, 5), (2, 3, 7, 9, 5)]
+    ```
+    Read the tutorial for introduction and application scenarios.
+    """
+    n_axes_before, n_axes_after, min_axes = analyze_pattern(pattern, "pack")
+    # packing zero tensors is illegal
+    backend = get_backend(tensors[0])
+    reshaped_tensors: List[Tensor] = []
+    packed_shapes: List[Shape] = []
+    for i, tensor in enumerate(tensors):
+        shape = backend.shape(tensor)
+        if len(shape) < min_axes:
+            raise EinopsError(
+                f"packed tensor #{i} (enumeration starts with 0) has shape {shape}, "
+                f"while pattern {pattern} assumes at least {min_axes} axes"
+            )
+        axis_after_packed_axes = len(shape) - n_axes_after
+        packed_shapes.append(shape[n_axes_before:axis_after_packed_axes])
+        reshaped_tensors.append(backend.reshape(tensor, (*shape[:n_axes_before], -1, *shape[axis_after_packed_axes:])))
+    return backend.concat(reshaped_tensors, axis=n_axes_before), packed_shapes
+def prod(x: Shape) -> int:
+    result = 1
+    for i in x:
+        result *= i
+    return result
+def unpack(tensor: Tensor, packed_shapes: List[Shape], pattern: str) -> List[Tensor]:
+    """
+    Unpacks a single tensor into several by splitting over a selected axes.
+    See einops tutorial for introduction into packing (and how it replaces stack and concatenation).
+    Parameters:
+        tensor: tensor to be unpacked
+        packed_shapes: packed_shapes (aka PS) is a list of shapes that take place of '*' in each output.
+            output will contain a single tensor for every provided shape
+        pattern: pattern that is shared for input and all outputs, e.g. "i j * k" or "batch seq *",
+            where * designates an axis to be unpacked
+    Returns:
+        list of tensors
+    If framework supports views, results are views to the original tensor.
+    Example:
+    ```python
+    >>> from numpy import zeros as Z
+    >>> inputs = [Z([2, 3, 5]), Z([2, 3, 7, 5]), Z([2, 3, 7, 9, 5])]
+    >>> packed, ps = pack(inputs, 'i j * k')
+    >>> packed.shape, ps
+    ((2, 3, 71, 5), [(), (7,), (7, 9)])
+    ```
+    In this example, axes were matched to: i=2, j=3, k=5 based on order (first, second, and last).
+    All other axes were 'packed' and concatenated.
+    PS (packed shapes) contains information about axes that were matched to '*' in every input.
+    Resulting tensor has as many elements as all inputs in total.
+    Packing can be reversed with unpack, which additionally needs PS (packed shapes) to reconstruct order.
+    ```python
+    >>> inputs_unpacked = unpack(packed, ps, 'i j * k')
+    >>> [x.shape for x in inputs_unpacked]
+    [(2, 3, 5), (2, 3, 7, 5), (2, 3, 7, 9, 5)]
+    ```
+    Read the tutorial for introduction and application scenarios.
+    """
+    n_axes_before, n_axes_after, min_axes = analyze_pattern(pattern, opname="unpack")
+    backend = get_backend(tensor)
+    input_shape = backend.shape(tensor)
+    if len(input_shape) != n_axes_before + 1 + n_axes_after:
+        raise EinopsError(f"unpack(..., {pattern}) received input of wrong dim with shape {input_shape}")
+    unpacked_axis: int = n_axes_before
+    lengths_of_composed_axes: List[int] = [-1 if -1 in p_shape else prod(p_shape) for p_shape in packed_shapes]
+    n_unknown_composed_axes = sum(int(x == -1) for x in lengths_of_composed_axes)
+    if n_unknown_composed_axes > 1:
+        raise EinopsError(
+            f"unpack(..., {pattern}) received more than one -1 in {packed_shapes} and can't infer dimensions"
+        )
+    # following manipulations allow to skip some shape verifications
+    # and leave it to backends
+    # [[], [2, 3], [4], [-1, 5], [6]] < examples of packed_axis
+    # split positions when computed should be
+    # [0,   1,      7,   11,      N-6 , N ], where N = length of axis
+    split_positions = [0] * len(packed_shapes) + [input_shape[unpacked_axis]]
+    if n_unknown_composed_axes == 0:
+        for i, x in enumerate(lengths_of_composed_axes[:-1]):
+            split_positions[i + 1] = split_positions[i] + x
+    else:
+        unknown_composed_axis: int = lengths_of_composed_axes.index(-1)
+        for i in range(unknown_composed_axis):
+            split_positions[i + 1] = split_positions[i] + lengths_of_composed_axes[i]
+        for j in range(unknown_composed_axis + 1, len(lengths_of_composed_axes))[::-1]:
+            split_positions[j] = split_positions[j + 1] - lengths_of_composed_axes[j]
+    shape_start = input_shape[:unpacked_axis]
+    shape_end = input_shape[unpacked_axis + 1 :]
+    slice_filler = (slice(None, None),) * unpacked_axis
+    try:
+        return [
+            backend.reshape(
+                # shortest way slice arbitrary axis
+                tensor[(*slice_filler, slice(split_positions[i], split_positions[i + 1]))],
+                (*shape_start, *element_shape, *shape_end),
+            )
+            for i, element_shape in enumerate(packed_shapes)
+        ]
+    except Exception:
+        # this hits if there is an error during reshapes, which means passed shapes were incorrect
+        raise RuntimeError(
+            f'Error during unpack(..., "{pattern}"): could not split axis of size {split_positions[-1]}'
+            f" into requested {packed_shapes}"
+        )

build/torch-cuda/parsing.py ADDED Viewed

	@@ -0,0 +1,152 @@

+from . import EinopsError
+import keyword
+import warnings
+from typing import List, Optional, Set, Tuple, Union
+_ellipsis: str = "…"  # NB, this is a single unicode symbol. String is used as it is not a list, but can be iterated
+class AnonymousAxis(object):
+    """Important thing: all instances of this class are not equal to each other"""
+    def __init__(self, value: str):
+        self.value = int(value)
+        if self.value <= 1:
+            if self.value == 1:
+                raise EinopsError("No need to create anonymous axis of length 1. Report this as an issue")
+            else:
+                raise EinopsError("Anonymous axis should have positive length, not {}".format(self.value))
+    def __repr__(self):
+        return "{}-axis".format(str(self.value))
+class ParsedExpression:
+    """
+    non-mutable structure that contains information about one side of expression (e.g. 'b c (h w)')
+    and keeps some information important for downstream
+    """
+    def __init__(self, expression: str, *, allow_underscore: bool = False, allow_duplicates: bool = False):
+        self.has_ellipsis: bool = False
+        self.has_ellipsis_parenthesized: Optional[bool] = None
+        self.identifiers: Set[str] = set()
+        # that's axes like 2, 3, 4 or 5. Axes with size 1 are exceptional and replaced with empty composition
+        self.has_non_unitary_anonymous_axes: bool = False
+        # composition keeps structure of composite axes, see how different corner cases are handled in tests
+        self.composition: List[Union[List[str], str]] = []
+        if "." in expression:
+            if "..." not in expression:
+                raise EinopsError("Expression may contain dots only inside ellipsis (...)")
+            if str.count(expression, "...") != 1 or str.count(expression, ".") != 3:
+                raise EinopsError(
+                    "Expression may contain dots only inside ellipsis (...); only one ellipsis for tensor "
+                )
+            expression = expression.replace("...", _ellipsis)
+            self.has_ellipsis = True
+        bracket_group: Optional[List[str]] = None
+        def add_axis_name(x):
+            if x in self.identifiers:
+                if not (allow_underscore and x == "_") and not allow_duplicates:
+                    raise EinopsError('Indexing expression contains duplicate dimension "{}"'.format(x))
+            if x == _ellipsis:
+                self.identifiers.add(_ellipsis)
+                if bracket_group is None:
+                    self.composition.append(_ellipsis)
+                    self.has_ellipsis_parenthesized = False
+                else:
+                    bracket_group.append(_ellipsis)
+                    self.has_ellipsis_parenthesized = True
+            else:
+                is_number = str.isdecimal(x)
+                if is_number and int(x) == 1:
+                    # handling the case of anonymous axis of length 1
+                    if bracket_group is None:
+                        self.composition.append([])
+                    else:
+                        pass  # no need to think about 1s inside parenthesis
+                    return
+                is_axis_name, reason = self.check_axis_name_return_reason(x, allow_underscore=allow_underscore)
+                if not (is_number or is_axis_name):
+                    raise EinopsError("Invalid axis identifier: {}\n{}".format(x, reason))
+                if is_number:
+                    x = AnonymousAxis(x)
+                self.identifiers.add(x)
+                if is_number:
+                    self.has_non_unitary_anonymous_axes = True
+                if bracket_group is None:
+                    self.composition.append([x])
+                else:
+                    bracket_group.append(x)
+        current_identifier = None
+        for char in expression:
+            if char in "() ":
+                if current_identifier is not None:
+                    add_axis_name(current_identifier)
+                current_identifier = None
+                if char == "(":
+                    if bracket_group is not None:
+                        raise EinopsError("Axis composition is one-level (brackets inside brackets not allowed)")
+                    bracket_group = []
+                elif char == ")":
+                    if bracket_group is None:
+                        raise EinopsError("Brackets are not balanced")
+                    self.composition.append(bracket_group)
+                    bracket_group = None
+            elif str.isalnum(char) or char in ["_", _ellipsis]:
+                if current_identifier is None:
+                    current_identifier = char
+                else:
+                    current_identifier += char
+            else:
+                raise EinopsError("Unknown character '{}'".format(char))
+        if bracket_group is not None:
+            raise EinopsError('Imbalanced parentheses in expression: "{}"'.format(expression))
+        if current_identifier is not None:
+            add_axis_name(current_identifier)
+    def flat_axes_order(self) -> List:
+        result = []
+        for composed_axis in self.composition:
+            assert isinstance(composed_axis, list), "does not work with ellipsis"
+            for axis in composed_axis:
+                result.append(axis)
+        return result
+    def has_composed_axes(self) -> bool:
+        # this will ignore 1 inside brackets
+        for axes in self.composition:
+            if isinstance(axes, list) and len(axes) > 1:
+                return True
+        return False
+    @staticmethod
+    def check_axis_name_return_reason(name: str, allow_underscore: bool = False) -> Tuple[bool, str]:
+        if not str.isidentifier(name):
+            return False, "not a valid python identifier"
+        elif name[0] == "_" or name[-1] == "_":
+            if name == "_" and allow_underscore:
+                return True, ""
+            return False, "axis name should should not start or end with underscore"
+        else:
+            if keyword.iskeyword(name):
+                warnings.warn("It is discouraged to use axes names that are keywords: {}".format(name), RuntimeWarning)
+            if name in ["axis"]:
+                warnings.warn(
+                    "It is discouraged to use 'axis' as an axis name " "and will raise an error in future",
+                    FutureWarning,
+                )
+            return True, ""
+    @staticmethod
+    def check_axis_name(name: str) -> bool:
+        """
+        Valid axes names are python identifiers except keywords,
+        and additionally should not start or end with underscore
+        """
+        is_valid, _reason = ParsedExpression.check_axis_name_return_reason(name)
+        return is_valid