Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

kronos-venv/lib/python3.12/site-packages/einops/__init__.py

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

kronos-venv/lib/python3.12/site-packages/einops/_backends.py

@@ -0,0 +1,764 @@
+"""
+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
+
+    previously_loaded_backends = list(_loaded_backends.items())
+    for _framework_name, backend in previously_loaded_backends:
+        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)
+
+    # handles modification of _loaded_backends from other thread, see #391
+    prev_backend_names = [x for x, _ in previously_loaded_backends]
+    for BackendSubclass in backend_subclasses:
+        if _debug_importing:
+            print("Testing for subclass of ", BackendSubclass)
+        if BackendSubclass.framework_name not in prev_backend_names:
+            # 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(f"Tensor type unknown to einops {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 f"<einops backend for {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().__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 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):
+        yield from self.elements
+
+    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)
+
+
+class MLXBackend(AbstractBackend):
+    framework_name = "mlx"
+
+    def __init__(self):
+        import mlx.core as mx
+        import numpy as np
+
+        self.mx = mx
+        self.np = np
+
+    def is_appropriate_type(self, tensor):
+        return isinstance(tensor, self.mx.array)
+
+    def from_numpy(self, x):
+        return self.mx.array(x)
+
+    def to_numpy(self, x):
+        if x.dtype == self.mx.bfloat16:
+            x = x.astype(self.mx.float32)
+        return self.np.array(x)
+
+    def arange(self, start, stop):
+        return self.mx.arange(start, stop)
+
+    def stack_on_zeroth_dimension(self, tensors: list):
+        return self.mx.stack(tensors)
+
+    def add_axes(self, x, new_position):
+        return self.mx.expand_dims(x, new_position)
+
+    def tile(self, x, repeats):
+        return self.mx.tile(x, repeats)
+
+    def concat(self, tensors, axis: int):
+        return self.mx.concatenate(tensors, axis=axis)
+
+    def is_float_type(self, x):
+        return self.mx.issubdtype(x.dtype, self.mx.floating)
+
+    def einsum(self, pattern, *x):
+        return self.mx.einsum(pattern, *x)

kronos-venv/lib/python3.12/site-packages/einops/_torch_specific.py

@@ -0,0 +1,136 @@
+"""
+Specialization of einops for torch.
+
+Unfortunately, torch's jit scripting mechanism isn't strong enough,
+and to have scripting supported at least for layers,
+a number of additional moves is needed.
+
+Design of main operations (dynamic resolution by lookup) is unlikely
+to be implemented by torch.jit.script,
+but torch.compile seems to work with operations just fine.
+"""
+
+import warnings
+from typing import Dict, List, Tuple
+
+import torch
+
+from einops.einops import TransformRecipe, _reconstruct_from_shape_uncached
+
+
+class TorchJitBackend:
+    """
+    Completely static backend that mimics part of normal backend functionality
+    but restricted to be within torchscript.
+    """
+
+    @staticmethod
+    def reduce(x: torch.Tensor, operation: str, reduced_axes: List[int]):
+        if operation == "min":
+            return x.amin(dim=reduced_axes)
+        elif operation == "max":
+            return x.amax(dim=reduced_axes)
+        elif operation == "sum":
+            return x.sum(dim=reduced_axes)
+        elif operation == "mean":
+            return x.mean(dim=reduced_axes)
+        elif operation == "prod":
+            for i in 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
+
+    if hasattr(torch, "__version__") and torch.__version__ >= "2.8":
+        # einops don't need to use allow_in graph for torch 2.8 and above
+        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, stacklevel=1
+        )
+        return
+
+    from .einops import einsum, rearrange, reduce, repeat
+    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()

kronos-venv/lib/python3.12/site-packages/einops/array_api.py

@@ -0,0 +1,125 @@
+from typing import List, Sequence, Tuple
+
+from .einops import EinopsError, Reduction, Tensor, _apply_recipe_array_api, _prepare_transformation_recipe
+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 = f' Error while processing {reduction}-reduction pattern "{pattern}".'
+        if not isinstance(tensor, list):
+            message += f"\n Input tensor shape: {tensor.shape}. "
+        else:
+            message += "\n Input is list. "
+        message += f"Additional info: {axes_lengths}."
+        raise EinopsError(message + f"\n {e}") from None
+
+
+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 as e:
+        # 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}"
+        ) from e

kronos-venv/lib/python3.12/site-packages/einops/einops.py

@@ -0,0 +1,939 @@
+import functools
+import itertools
+import string
+import typing
+from collections import OrderedDict
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, TypeVar, Union, cast, overload
+
+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 AnonymousAxis, ParsedExpression, _ellipsis
+
+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(f"Ellipsis found in right side, but not left side of a pattern {pattern}")
+    if left.has_ellipsis and left.has_ellipsis_parenthesized:
+        raise EinopsError(f"Ellipsis inside parenthesis in the left side is not allowed: {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(f"Identifiers only on one side of expression (should be on both): {difference}")
+    elif operation == "repeat":
+        difference = set.difference(left.identifiers, rght.identifiers)
+        if len(difference) > 0:
+            raise EinopsError(f"Unexpected identifiers on the left side of repeat: {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(f"Specify sizes for new axes in repeat: {axes_without_size}")
+    elif operation in _reductions or callable(operation):
+        difference = set.difference(rght.identifiers, left.identifiers)
+        if len(difference) > 0:
+            raise EinopsError(f"Unexpected identifiers on the right side of reduce {operation}: {difference}")
+    else:
+        raise EinopsError(f"Unknown reduction {operation}. Expect one of {_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(f"Axis {elementary_axis} is not used in transform")
+        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 composite_axis in 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(f"Could not infer sizes for {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}
+
+
+@overload
+def reduce(tensor: List[Tensor], pattern: str, reduction: Reduction, **axes_lengths: Size) -> Tensor: ...
+
+
+@overload
+def reduce(tensor: Tensor, pattern: str, reduction: Reduction, **axes_lengths: Size) -> Tensor: ...
+
+
+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 = f' Error while processing {reduction}-reduction pattern "{pattern}".'
+        if not isinstance(tensor, list):
+            message += f"\n Input tensor shape: {shape}. "
+        else:
+            message += "\n Input is list. "
+        message += f"Additional info: {axes_lengths}."
+        raise EinopsError(message + f"\n {e}") from None
+
+
+@overload
+def rearrange(tensor: List[Tensor], pattern: str, **axes_lengths: Size) -> Tensor: ...
+
+
+@overload
+def rearrange(tensor: Tensor, pattern: str, **axes_lengths: Size) -> Tensor: ...
+
+
+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)
+
+
+@overload
+def repeat(tensor: List[Tensor], pattern: str, **axes_lengths: Size) -> Tensor: ...
+
+
+@overload
+def repeat(tensor: Tensor, pattern: str, **axes_lengths: Size) -> Tensor: ...
+
+
+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, i.e. upsample an image by 2x
+    >>> repeat(image, 'h w -> (h h2) (w w2)', h2=2, w2=2).shape
+    (60, 80)
+
+    # 'pixelate' an 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)

kronos-venv/lib/python3.12/site-packages/einops/experimental/indexing.py

@@ -0,0 +1,5 @@
+"""
+This file contained some thoughts on indexing.
+
+These ideas were developed further in eindex (separate package).
+"""

kronos-venv/lib/python3.12/site-packages/einops/layers/__init__.py

@@ -0,0 +1,105 @@
+__author__ = "Alex Rogozhnikov"
+
+from typing import Any, Dict
+
+from einops import EinopsError
+from einops.einops import TransformRecipe, _apply_recipe, _prepare_recipes_for_all_dims, get_backend
+
+
+class RearrangeMixin:
+    """
+    Rearrange layer behaves identically to einops.rearrange operation.
+
+    :param pattern: str, rearrangement pattern
+    :param axes_lengths: any additional specification of dimensions
+
+    See einops.rearrange for source_examples.
+    """
+
+    def __init__(self, pattern: str, **axes_lengths: Any) -> None:
+        super().__init__()
+        self.pattern = pattern
+        self.axes_lengths = axes_lengths
+        # self._recipe = self.recipe()  # checking parameters
+        self._multirecipe = self.multirecipe()
+        self._axes_lengths = tuple(self.axes_lengths.items())
+
+    def __repr__(self) -> str:
+        params = repr(self.pattern)
+        for axis, length in self.axes_lengths.items():
+            params += f", {axis}={length}"
+        return f"{self.__class__.__name__}({params})"
+
+    def multirecipe(self) -> Dict[int, TransformRecipe]:
+        try:
+            return _prepare_recipes_for_all_dims(
+                self.pattern, operation="rearrange", axes_names=tuple(self.axes_lengths)
+            )
+        except EinopsError as e:
+            raise EinopsError(f" Error while preparing {self!r}\n {e}") from None
+
+    def _apply_recipe(self, x):
+        backend = get_backend(x)
+        return _apply_recipe(
+            backend=backend,
+            recipe=self._multirecipe[len(x.shape)],
+            tensor=x,
+            reduction_type="rearrange",
+            axes_lengths=self._axes_lengths,
+        )
+
+    def __getstate__(self):
+        return {"pattern": self.pattern, "axes_lengths": self.axes_lengths}
+
+    def __setstate__(self, state):
+        self.__init__(pattern=state["pattern"], **state["axes_lengths"])
+
+
+class ReduceMixin:
+    """
+    Reduce layer behaves identically to einops.reduce operation.
+
+    :param pattern: str, rearrangement pattern
+    :param reduction: one of available reductions ('min', 'max', 'sum', 'mean', 'prod'), case-sensitive
+    :param axes_lengths: any additional specification of dimensions
+
+    See einops.reduce for source_examples.
+    """
+
+    def __init__(self, pattern: str, reduction: str, **axes_lengths: Any):
+        super().__init__()
+        self.pattern = pattern
+        self.reduction = reduction
+        self.axes_lengths = axes_lengths
+        self._multirecipe = self.multirecipe()
+        self._axes_lengths = tuple(self.axes_lengths.items())
+
+    def __repr__(self):
+        params = f"{self.pattern!r}, {self.reduction!r}"
+        for axis, length in self.axes_lengths.items():
+            params += f", {axis}={length}"
+        return f"{self.__class__.__name__}({params})"
+
+    def multirecipe(self) -> Dict[int, TransformRecipe]:
+        try:
+            return _prepare_recipes_for_all_dims(
+                self.pattern, operation=self.reduction, axes_names=tuple(self.axes_lengths)
+            )
+        except EinopsError as e:
+            raise EinopsError(f" Error while preparing {self!r}\n {e}") from None
+
+    def _apply_recipe(self, x):
+        backend = get_backend(x)
+        return _apply_recipe(
+            backend=backend,
+            recipe=self._multirecipe[len(x.shape)],
+            tensor=x,
+            reduction_type=self.reduction,
+            axes_lengths=self._axes_lengths,
+        )
+
+    def __getstate__(self):
+        return {"pattern": self.pattern, "reduction": self.reduction, "axes_lengths": self.axes_lengths}
+
+    def __setstate__(self, state):
+        self.__init__(pattern=state["pattern"], reduction=state["reduction"], **state["axes_lengths"])

kronos-venv/lib/python3.12/site-packages/einops/layers/_einmix.py

@@ -0,0 +1,227 @@
+import string
+import warnings
+from typing import Any, Dict, List, Optional
+
+from einops import EinopsError
+from einops.einops import _product
+from einops.parsing import ParsedExpression, _ellipsis
+
+
+def _report_axes(axes: set, report_message: str):
+    if len(axes) > 0:
+        raise EinopsError(report_message.format(axes))
+
+
+class _EinmixMixin:
+    def __init__(self, pattern: str, weight_shape: str, bias_shape: Optional[str] = None, **axes_lengths: Any):
+        """
+        EinMix - Einstein summation with automated tensor management and axis packing/unpacking.
+
+        EinMix is a combination of einops and MLP, see tutorial:
+        https://github.com/arogozhnikov/einops/blob/main/docs/3-einmix-layer.ipynb
+
+        Imagine taking einsum with two arguments, one of each input, and one - tensor with weights
+        >>> einsum('time batch channel_in, channel_in channel_out -> time batch channel_out', input, weight)
+
+        This layer manages weights for you, syntax highlights a special role of weight matrix
+        >>> EinMix('time batch channel_in -> time batch channel_out', weight_shape='channel_in channel_out')
+        But otherwise it is the same einsum under the hood. Plus einops-rearrange.
+
+        Simple linear layer with a bias term (you have one like that in your framework)
+        >>> EinMix('t b cin -> t b cout', weight_shape='cin cout', bias_shape='cout', cin=10, cout=20)
+        There is no restriction to mix the last axis. Let's mix along height
+        >>> EinMix('h w c-> hout w c', weight_shape='h hout', bias_shape='hout', h=32, hout=32)
+        Example of channel-wise multiplication (like one used in normalizations)
+        >>> EinMix('t b c -> t b c', weight_shape='c', c=128)
+        Multi-head linear layer (each head is own linear layer):
+        >>> EinMix('t b (head cin) -> t b (head cout)', weight_shape='head cin cout', ...)
+
+        ... and yes, you need to specify all dimensions of weight shape/bias shape in parameters.
+
+        Use cases:
+        - when channel dimension is not last, use EinMix, not transposition
+        - patch/segment embeddings
+        - when need only within-group connections to reduce number of weights and computations
+        - next-gen MLPs (follow tutorial link above to learn more!)
+        - in general, any time you want to combine linear layer and einops.rearrange
+
+        Uniform He initialization is applied to weight tensor.
+        This accounts for the number of elements mixed and produced.
+
+        Parameters
+        :param pattern: transformation pattern, left side - dimensions of input, right side - dimensions of output
+        :param weight_shape: axes of weight. A tensor of this shape is created, stored, and optimized in a layer
+               If bias_shape is not specified, bias is not created.
+        :param bias_shape: axes of bias added to output. Weights of this shape are created and stored. If `None` (the default), no bias is added.
+        :param axes_lengths: dimensions of weight tensor
+        """
+        super().__init__()
+        self.pattern = pattern
+        self.weight_shape = weight_shape
+        self.bias_shape = bias_shape
+        self.axes_lengths = axes_lengths
+        self.initialize_einmix(
+            pattern=pattern, weight_shape=weight_shape, bias_shape=bias_shape, axes_lengths=axes_lengths
+        )
+
+    def initialize_einmix(self, pattern: str, weight_shape: str, bias_shape: Optional[str], axes_lengths: dict):
+        left_pattern, right_pattern = pattern.split("->")
+        left = ParsedExpression(left_pattern)
+        right = ParsedExpression(right_pattern)
+        weight = ParsedExpression(weight_shape)
+        _report_axes(
+            set.difference(right.identifiers, {*left.identifiers, *weight.identifiers}),
+            "Unrecognized identifiers on the right side of EinMix {}",
+        )
+        if weight.has_ellipsis:
+            raise EinopsError("Ellipsis is not supported in weight, as its shape should be fully specified")
+        if left.has_ellipsis or right.has_ellipsis:
+            if not (left.has_ellipsis and right.has_ellipsis):
+                raise EinopsError(f"Ellipsis in EinMix should be on both sides, {pattern}")
+            if left.has_ellipsis_parenthesized:
+                raise EinopsError(f"Ellipsis on left side can't be in parenthesis, got {pattern}")
+        if any(x.has_non_unitary_anonymous_axes for x in [left, right, weight]):
+            raise EinopsError("Anonymous axes (numbers) are not allowed in EinMix")
+        if "(" in weight_shape or ")" in weight_shape:
+            raise EinopsError(f"Parenthesis is not allowed in weight shape: {weight_shape}")
+
+        pre_reshape_pattern = None
+        pre_reshape_lengths = None
+        post_reshape_pattern = None
+        if any(len(group) != 1 for group in left.composition):
+            names: List[str] = []
+            for group in left.composition:
+                names += group
+            names = [name if name != _ellipsis else "..." for name in names]
+            composition = " ".join(names)
+            pre_reshape_pattern = f"{left_pattern}-> {composition}"
+            pre_reshape_lengths = {name: length for name, length in axes_lengths.items() if name in names}
+
+        if any(len(group) != 1 for group in right.composition) or right.has_ellipsis_parenthesized:
+            names = []
+            for group in right.composition:
+                names += group
+            names = [name if name != _ellipsis else "..." for name in names]
+            composition = " ".join(names)
+            post_reshape_pattern = f"{composition} ->{right_pattern}"
+
+        self._create_rearrange_layers(pre_reshape_pattern, pre_reshape_lengths, post_reshape_pattern, {})
+
+        for axis in weight.identifiers:
+            if axis not in axes_lengths:
+                raise EinopsError(f"Dimension {axis} of weight should be specified")
+        _report_axes(
+            set.difference(set(axes_lengths), {*left.identifiers, *weight.identifiers}),
+            "Axes {} are not used in pattern",
+        )
+        _report_axes(
+            set.difference(weight.identifiers, {*left.identifiers, *right.identifiers}), "Weight axes {} are redundant"
+        )
+        if len(weight.identifiers) == 0:
+            warnings.warn("EinMix: weight has no dimensions (means multiplication by a number)", stacklevel=2)
+
+        _weight_shape = [axes_lengths[axis] for (axis,) in weight.composition]
+        # single output element is a combination of fan_in input elements
+        _fan_in = _product([axes_lengths[axis] for (axis,) in weight.composition if axis not in right.identifiers])
+        if bias_shape is not None:
+            # maybe I should put ellipsis in the beginning for simplicity?
+            if not isinstance(bias_shape, str):
+                raise EinopsError("bias shape should be string specifying which axes bias depends on")
+            bias = ParsedExpression(bias_shape)
+            _report_axes(
+                set.difference(bias.identifiers, right.identifiers),
+                "Bias axes {} not present in output",
+            )
+            _report_axes(
+                set.difference(bias.identifiers, set(axes_lengths)),
+                "Sizes not provided for bias axes {}",
+            )
+
+            _bias_shape = []
+            used_non_trivial_size = False
+            for axes in right.composition:
+                if axes == _ellipsis:
+                    if used_non_trivial_size:
+                        raise EinopsError("all bias dimensions should go after ellipsis in the output")
+                else:
+                    # handles ellipsis correctly
+                    for axis in axes:
+                        if axis == _ellipsis:
+                            if used_non_trivial_size:
+                                raise EinopsError("all bias dimensions should go after ellipsis in the output")
+                        elif axis in bias.identifiers:
+                            _bias_shape.append(axes_lengths[axis])
+                            used_non_trivial_size = True
+                        else:
+                            _bias_shape.append(1)
+        else:
+            _bias_shape = None
+
+        weight_bound = (3 / _fan_in) ** 0.5
+        bias_bound = (1 / _fan_in) ** 0.5
+        self._create_parameters(_weight_shape, weight_bound, _bias_shape, bias_bound)
+
+        # rewrite einsum expression with single-letter latin identifiers so that
+        # expression will be understood by any framework
+        mapped_identifiers = {*left.identifiers, *right.identifiers, *weight.identifiers}
+        if _ellipsis in mapped_identifiers:
+            mapped_identifiers.remove(_ellipsis)
+        mapped_identifiers = sorted(mapped_identifiers)
+        mapping2letters = {k: letter for letter, k in zip(string.ascii_lowercase, mapped_identifiers)}
+        mapping2letters[_ellipsis] = "..."  # preserve ellipsis
+
+        def write_flat_remapped(axes: ParsedExpression):
+            result = []
+            for composed_axis in axes.composition:
+                if isinstance(composed_axis, list):
+                    result.extend([mapping2letters[axis] for axis in composed_axis])
+                else:
+                    assert composed_axis == _ellipsis
+                    result.append("...")
+            return "".join(result)
+
+        self.einsum_pattern: str = (
+            f"{write_flat_remapped(left)},{write_flat_remapped(weight)}->{write_flat_remapped(right)}"
+        )
+
+    def _create_rearrange_layers(
+        self,
+        pre_reshape_pattern: Optional[str],
+        pre_reshape_lengths: Optional[Dict],
+        post_reshape_pattern: Optional[str],
+        post_reshape_lengths: Optional[Dict],
+    ):
+        raise NotImplementedError("Should be defined in framework implementations")
+
+    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
+        """Shape and implementations"""
+        raise NotImplementedError("Should be defined in framework implementations")
+
+    def __repr__(self):
+        params = repr(self.pattern)
+        params += f", '{self.weight_shape}'"
+        if self.bias_shape is not None:
+            params += f", '{self.bias_shape}'"
+        for axis, length in self.axes_lengths.items():
+            params += f", {axis}={length}"
+        return f"{self.__class__.__name__}({params})"
+
+
+class _EinmixDebugger(_EinmixMixin):
+    """Used only to test mixin"""
+
+    def _create_rearrange_layers(
+        self,
+        pre_reshape_pattern: Optional[str],
+        pre_reshape_lengths: Optional[Dict],
+        post_reshape_pattern: Optional[str],
+        post_reshape_lengths: Optional[Dict],
+    ):
+        self.pre_reshape_pattern = pre_reshape_pattern
+        self.pre_reshape_lengths = pre_reshape_lengths
+        self.post_reshape_pattern = post_reshape_pattern
+        self.post_reshape_lengths = post_reshape_lengths
+
+    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
+        self.saved_weight_shape = weight_shape
+        self.saved_bias_shape = bias_shape

kronos-venv/lib/python3.12/site-packages/einops/layers/flax.py

@@ -0,0 +1,82 @@
+from dataclasses import field
+from typing import Dict, Optional, cast
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+
+from . import RearrangeMixin, ReduceMixin
+from ._einmix import _EinmixMixin
+
+__author__ = "Alex Rogozhnikov"
+
+
+class Reduce(nn.Module):
+    pattern: str
+    reduction: str
+    sizes: dict = field(default_factory=dict)
+
+    def setup(self):
+        self.reducer = ReduceMixin(self.pattern, self.reduction, **self.sizes)
+
+    def __call__(self, input):
+        return self.reducer._apply_recipe(input)
+
+
+class Rearrange(nn.Module):
+    pattern: str
+    sizes: dict = field(default_factory=dict)
+
+    def setup(self):
+        self.rearranger = RearrangeMixin(self.pattern, **self.sizes)
+
+    def __call__(self, input):
+        return self.rearranger._apply_recipe(input)
+
+
+class EinMix(nn.Module, _EinmixMixin):
+    pattern: str
+    weight_shape: str
+    bias_shape: Optional[str] = None
+    sizes: dict = field(default_factory=dict)
+
+    def setup(self):
+        self.initialize_einmix(
+            pattern=self.pattern,
+            weight_shape=self.weight_shape,
+            bias_shape=self.bias_shape,
+            axes_lengths=self.sizes,
+        )
+
+    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
+        self.weight = self.param("weight", jax.nn.initializers.uniform(weight_bound), weight_shape)
+
+        if bias_shape is not None:
+            self.bias = self.param("bias", jax.nn.initializers.uniform(bias_bound), bias_shape)
+        else:
+            self.bias = None
+
+    def _create_rearrange_layers(
+        self,
+        pre_reshape_pattern: Optional[str],
+        pre_reshape_lengths: Optional[Dict],
+        post_reshape_pattern: Optional[str],
+        post_reshape_lengths: Optional[Dict],
+    ):
+        self.pre_rearrange = None
+        if pre_reshape_pattern is not None:
+            self.pre_rearrange = Rearrange(pre_reshape_pattern, sizes=cast(dict, pre_reshape_lengths))
+
+        self.post_rearrange = None
+        if post_reshape_pattern is not None:
+            self.post_rearrange = Rearrange(post_reshape_pattern, sizes=cast(dict, post_reshape_lengths))
+
+    def __call__(self, input):
+        if self.pre_rearrange is not None:
+            input = self.pre_rearrange(input)
+        result = jnp.einsum(self.einsum_pattern, input, self.weight)
+        if self.bias is not None:
+            result += self.bias
+        if self.post_rearrange is not None:
+            result = self.post_rearrange(result)
+        return result

kronos-venv/lib/python3.12/site-packages/einops/layers/keras.py

@@ -0,0 +1,9 @@
+__author__ = "Alex Rogozhnikov"
+
+from einops.layers.tensorflow import EinMix, Rearrange, Reduce
+
+keras_custom_objects = {
+    Rearrange.__name__: Rearrange,
+    Reduce.__name__: Reduce,
+    EinMix.__name__: EinMix,
+}

kronos-venv/lib/python3.12/site-packages/einops/layers/oneflow.py

@@ -0,0 +1,54 @@
+from typing import Dict, Optional, cast
+
+import oneflow as flow
+
+from . import RearrangeMixin, ReduceMixin
+from ._einmix import _EinmixMixin
+
+__author__ = "Tianhe Ren & Depeng Liang"
+
+
+class Rearrange(RearrangeMixin, flow.nn.Module):
+    def forward(self, input):
+        return self._apply_recipe(input)
+
+
+class Reduce(ReduceMixin, flow.nn.Module):
+    def forward(self, input):
+        return self._apply_recipe(input)
+
+
+class EinMix(_EinmixMixin, flow.nn.Module):
+    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
+        self.weight = flow.nn.Parameter(
+            flow.zeros(weight_shape).uniform_(-weight_bound, weight_bound), requires_grad=True
+        )
+        if bias_shape is not None:
+            self.bias = flow.nn.Parameter(flow.zeros(bias_shape).uniform_(-bias_bound, bias_bound), requires_grad=True)
+        else:
+            self.bias = None
+
+    def _create_rearrange_layers(
+        self,
+        pre_reshape_pattern: Optional[str],
+        pre_reshape_lengths: Optional[Dict],
+        post_reshape_pattern: Optional[str],
+        post_reshape_lengths: Optional[Dict],
+    ):
+        self.pre_rearrange = None
+        if pre_reshape_pattern is not None:
+            self.pre_rearrange = Rearrange(pre_reshape_pattern, **cast(dict, pre_reshape_lengths))
+
+        self.post_rearrange = None
+        if post_reshape_pattern is not None:
+            self.post_rearrange = Rearrange(post_reshape_pattern, **cast(dict, post_reshape_lengths))
+
+    def forward(self, input):
+        if self.pre_rearrange is not None:
+            input = self.pre_rearrange(input)
+        result = flow.einsum(self.einsum_pattern, input, self.weight)
+        if self.bias is not None:
+            result += self.bias
+        if self.post_rearrange is not None:
+            result = self.post_rearrange(result)
+        return result

kronos-venv/lib/python3.12/site-packages/einops/layers/paddle.py

@@ -0,0 +1,58 @@
+from typing import Dict, Optional, cast
+
+import paddle
+
+from . import RearrangeMixin, ReduceMixin
+from ._einmix import _EinmixMixin
+
+__author__ = "PaddlePaddle"
+
+
+class Rearrange(RearrangeMixin, paddle.nn.Layer):
+    def forward(self, input):
+        return self._apply_recipe(input)
+
+
+class Reduce(ReduceMixin, paddle.nn.Layer):
+    def forward(self, input):
+        return self._apply_recipe(input)
+
+
+class EinMix(_EinmixMixin, paddle.nn.Layer):
+    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
+        self.weight = self.create_parameter(
+            weight_shape, default_initializer=paddle.nn.initializer.Uniform(-weight_bound, weight_bound)
+        )
+
+        if bias_shape is not None:
+            self.bias = self.create_parameter(
+                bias_shape, default_initializer=paddle.nn.initializer.Uniform(-bias_bound, bias_bound)
+            )
+        else:
+            self.bias = None
+
+    def _create_rearrange_layers(
+        self,
+        pre_reshape_pattern: Optional[str],
+        pre_reshape_lengths: Optional[Dict],
+        post_reshape_pattern: Optional[str],
+        post_reshape_lengths: Optional[Dict],
+    ):
+        self.pre_rearrange = None
+        if pre_reshape_pattern is not None:
+            self.pre_rearrange = Rearrange(pre_reshape_pattern, **cast(dict, pre_reshape_lengths))
+
+        self.post_rearrange = None
+        if post_reshape_pattern is not None:
+            self.post_rearrange = Rearrange(post_reshape_pattern, **cast(dict, post_reshape_lengths))
+
+    def forward(self, input):
+        if self.pre_rearrange is not None:
+            input = self.pre_rearrange(input)
+
+        result = paddle.einsum(self.einsum_pattern, input, self.weight)
+        if self.bias is not None:
+            result += self.bias
+        if self.post_rearrange is not None:
+            result = self.post_rearrange(result)
+        return result

kronos-venv/lib/python3.12/site-packages/einops/layers/tensorflow.py

@@ -0,0 +1,101 @@
+"""
+Comment about tensorflow layers:
+unfortunately instructions on creation of TF layers change constantly,
+and changed way too many times at this point to remember what-compatible-where.
+
+Layers in einops==0.7.0 (and several prior versions)
+ are compatible with TF 2.13
+
+Layers in einops==0.8.0 were re-implemented
+ according to official instructions for TF 2.16
+
+"""
+
+from typing import Dict, Optional, cast
+
+import tensorflow as tf
+from tensorflow.keras.layers import Layer
+
+from . import RearrangeMixin, ReduceMixin
+from ._einmix import _EinmixMixin
+
+__author__ = "Alex Rogozhnikov"
+
+
+class Rearrange(RearrangeMixin, Layer):
+    def build(self, input_shape):
+        pass  # layer does not have any parameters to be initialized
+
+    def call(self, inputs):
+        return self._apply_recipe(inputs)
+
+    def get_config(self):
+        return {"pattern": self.pattern, **self.axes_lengths}
+
+
+class Reduce(ReduceMixin, Layer):
+    def build(self, input_shape):
+        pass  # layer does not have any parameters to be initialized
+
+    def call(self, inputs):
+        return self._apply_recipe(inputs)
+
+    def get_config(self):
+        return {"pattern": self.pattern, "reduction": self.reduction, **self.axes_lengths}
+
+
+class EinMix(_EinmixMixin, Layer):
+    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
+        # this method is called in __init__,
+        #  but we postpone actual creation to build(), as TF instruction suggests
+        self._params = [weight_shape, weight_bound, bias_shape, bias_bound]
+
+    def _create_rearrange_layers(
+        self,
+        pre_reshape_pattern: Optional[str],
+        pre_reshape_lengths: Optional[Dict],
+        post_reshape_pattern: Optional[str],
+        post_reshape_lengths: Optional[Dict],
+    ):
+        self.pre_rearrange = None
+        if pre_reshape_pattern is not None:
+            self.pre_rearrange = Rearrange(pre_reshape_pattern, **cast(dict, pre_reshape_lengths))
+
+        self.post_rearrange = None
+        if post_reshape_pattern is not None:
+            self.post_rearrange = Rearrange(post_reshape_pattern, **cast(dict, post_reshape_lengths))
+
+    def build(self, input_shape):
+        [weight_shape, weight_bound, bias_shape, bias_bound] = self._params
+        self.weight = self.add_weight(
+            shape=weight_shape,
+            initializer=tf.random_uniform_initializer(-weight_bound, weight_bound),
+            trainable=True,
+        )
+
+        if bias_shape is not None:
+            self.bias = self.add_weight(
+                shape=bias_shape,
+                initializer=tf.random_uniform_initializer(-bias_bound, bias_bound),
+                trainable=True,
+            )
+        else:
+            self.bias = None
+
+    def call(self, inputs):
+        if self.pre_rearrange is not None:
+            inputs = self.pre_rearrange(inputs)
+        result = tf.einsum(self.einsum_pattern, inputs, self.weight)
+        if self.bias is not None:
+            result = result + self.bias
+        if self.post_rearrange is not None:
+            result = self.post_rearrange(result)
+        return result
+
+    def get_config(self):
+        return {
+            "pattern": self.pattern,
+            "weight_shape": self.weight_shape,
+            "bias_shape": self.bias_shape,
+            **self.axes_lengths,
+        }

kronos-venv/lib/python3.12/site-packages/einops/layers/torch.py

@@ -0,0 +1,68 @@
+from typing import Dict, Optional, cast
+
+import torch
+
+from einops._torch_specific import apply_for_scriptable_torch
+
+from . import RearrangeMixin, ReduceMixin
+from ._einmix import _EinmixMixin
+
+__author__ = "Alex Rogozhnikov"
+
+
+class Rearrange(RearrangeMixin, torch.nn.Module):
+    def forward(self, input):
+        recipe = self._multirecipe[input.ndim]
+        return apply_for_scriptable_torch(recipe, input, reduction_type="rearrange", axes_dims=self._axes_lengths)
+
+    def _apply_recipe(self, x):
+        # overriding parent method to prevent it's scripting
+        pass
+
+
+class Reduce(ReduceMixin, torch.nn.Module):
+    def forward(self, input):
+        recipe = self._multirecipe[input.ndim]
+        return apply_for_scriptable_torch(recipe, input, reduction_type=self.reduction, axes_dims=self._axes_lengths)
+
+    def _apply_recipe(self, x):
+        # overriding parent method to prevent it's scripting
+        pass
+
+
+class EinMix(_EinmixMixin, torch.nn.Module):
+    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
+        self.weight = torch.nn.Parameter(
+            torch.zeros(weight_shape).uniform_(-weight_bound, weight_bound), requires_grad=True
+        )
+        if bias_shape is not None:
+            self.bias = torch.nn.Parameter(
+                torch.zeros(bias_shape).uniform_(-bias_bound, bias_bound), requires_grad=True
+            )
+        else:
+            self.bias = None
+
+    def _create_rearrange_layers(
+        self,
+        pre_reshape_pattern: Optional[str],
+        pre_reshape_lengths: Optional[Dict],
+        post_reshape_pattern: Optional[str],
+        post_reshape_lengths: Optional[Dict],
+    ):
+        self.pre_rearrange = None
+        if pre_reshape_pattern is not None:
+            self.pre_rearrange = Rearrange(pre_reshape_pattern, **cast(dict, pre_reshape_lengths))
+
+        self.post_rearrange = None
+        if post_reshape_pattern is not None:
+            self.post_rearrange = Rearrange(post_reshape_pattern, **cast(dict, post_reshape_lengths))
+
+    def forward(self, input):
+        if self.pre_rearrange is not None:
+            input = self.pre_rearrange(input)
+        result = torch.einsum(self.einsum_pattern, input, self.weight)
+        if self.bias is not None:
+            result += self.bias
+        if self.post_rearrange is not None:
+            result = self.post_rearrange(result)
+        return result

kronos-venv/lib/python3.12/site-packages/einops/packing.py

@@ -0,0 +1,189 @@
+from functools import lru_cache
+from typing import List, Sequence, Tuple, TypeVar, Union
+
+from einops import EinopsError
+from einops._backends import get_backend
+from einops.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 as e:
+        # this hits if there is an error during reshapes, which means passed shapes were incorrect
+        raise EinopsError(
+            f'Error during unpack(..., "{pattern}"): could not split axis of size {split_positions[-1]}'
+            f" into requested {packed_shapes}"
+        ) from e

kronos-venv/lib/python3.12/site-packages/einops/parsing.py

@@ -0,0 +1,158 @@
+import keyword
+import warnings
+from typing import List, Optional, Set, Tuple, Union
+
+from einops import EinopsError
+
+_ellipsis: str = "…"  # NB, this is a single unicode symbol. String is used as it is not a list, but can be iterated
+
+
+class AnonymousAxis:
+    """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(f"Anonymous axis should have positive length, not {self.value}")
+
+    def __repr__(self):
+        return f"{str(self.value)}-axis"
+
+
+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(f'Indexing expression contains duplicate dimension "{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(f"Invalid axis identifier: {x}\n{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(f"Unknown character '{char}'")
+
+        if bracket_group is not None:
+            raise EinopsError(f'Imbalanced parentheses in expression: "{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(
+                    f"It is discouraged to use axes names that are keywords: {name}",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+            if name in ["axis"]:
+                warnings.warn(
+                    "It is discouraged to use 'axis' as an axis name and will raise an error in future",
+                    FutureWarning,
+                    stacklevel=2,
+                )
+            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

kronos-venv/lib/python3.12/site-packages/einops/tests/__init__.py

@@ -0,0 +1,109 @@
+"""
+Common utils for testing.
+These functions allow testing only some frameworks, not all.
+"""
+
+import logging
+import os
+import warnings
+from functools import lru_cache
+from typing import List, Tuple
+
+from einops import _backends
+
+__author__ = "Alex Rogozhnikov"
+
+
+# minimize noise in tests logging
+logging.getLogger("tensorflow").disabled = True
+logging.getLogger("matplotlib").disabled = True
+
+FLOAT_REDUCTIONS = ("min", "max", "sum", "mean", "prod")  # not includes any/all
+
+
+def find_names_of_all_frameworks() -> List[str]:
+    backend_subclasses = []
+    backends = _backends.AbstractBackend.__subclasses__()
+    while backends:
+        backend = backends.pop()
+        backends += backend.__subclasses__()
+        backend_subclasses.append(backend)
+    return [b.framework_name for b in backend_subclasses]
+
+
+ENVVAR_NAME = "EINOPS_TEST_BACKENDS"
+
+
+def unparse_backends(backend_names: List[str]) -> Tuple[str, str]:
+    _known_backends = find_names_of_all_frameworks()
+    for backend_name in backend_names:
+        if backend_name not in _known_backends:
+            raise RuntimeError(f"Unknown framework: {backend_name}")
+    return ENVVAR_NAME, ",".join(backend_names)
+
+
+@lru_cache(maxsize=1)
+def parse_backends_to_test() -> List[str]:
+    if ENVVAR_NAME not in os.environ:
+        raise RuntimeError(f"Testing frameworks were not specified, env var {ENVVAR_NAME} not set")
+    parsed_backends = os.environ[ENVVAR_NAME].split(",")
+    _known_backends = find_names_of_all_frameworks()
+    for backend_name in parsed_backends:
+        if backend_name not in _known_backends:
+            raise RuntimeError(f"Unknown framework: {backend_name}")
+
+    return parsed_backends
+
+
+def is_backend_tested(backend: str) -> bool:
+    """Used to skip test if corresponding backend is not tested"""
+    if backend not in find_names_of_all_frameworks():
+        raise RuntimeError(f"Unknown framework {backend}")
+    return backend in parse_backends_to_test()
+
+
+def collect_test_backends(symbolic=False, layers=False) -> List[_backends.AbstractBackend]:
+    """
+    :param symbolic: symbolic or imperative frameworks?
+    :param layers: layers or operations?
+    :return: list of backends satisfying set conditions
+    """
+    if not symbolic:
+        if not layers:
+            backend_types = [
+                _backends.NumpyBackend,
+                _backends.JaxBackend,
+                _backends.TorchBackend,
+                _backends.TensorflowBackend,
+                _backends.OneFlowBackend,
+                _backends.PaddleBackend,
+                _backends.CupyBackend,
+            ]
+        else:
+            backend_types = [
+                _backends.TorchBackend,
+                _backends.OneFlowBackend,
+                _backends.PaddleBackend,
+            ]
+    else:
+        if not layers:
+            backend_types = [
+                _backends.PyTensorBackend,
+            ]
+        else:
+            backend_types = [
+                _backends.TFKerasBackend,
+            ]
+
+    backend_names_to_test = parse_backends_to_test()
+    result = []
+    for backend_type in backend_types:
+        if backend_type.framework_name not in backend_names_to_test:
+            continue
+        try:
+            result.append(backend_type())
+        except ImportError:
+            # problem with backend installation fails a specific test function,
+            # but will be skipped in all other test cases
+            warnings.warn(f"backend could not be initialized for tests: {backend_type}", stacklevel=1)
+    return result

kronos-venv/lib/python3.12/site-packages/einops/tests/run_tests.py

@@ -0,0 +1,89 @@
+"""
+Runs tests that are appropriate for framework.
+"""
+
+import os
+import sys
+from pathlib import Path
+from subprocess import Popen
+
+__author__ = "Alex Rogozhnikov"
+
+
+def run(cmd, **env):
+    # keeps printing output when testing
+    cmd = cmd.split(" ") if isinstance(cmd, str) else cmd
+    print("running:", cmd)
+    p = Popen(cmd, cwd=str(Path(__file__).parent), env={**os.environ, **env})
+    p.communicate()
+    return p.returncode
+
+
+def main():
+    _executable, *args = sys.argv
+    frameworks = [x for x in args if x != "--pip-install"]
+    pip_install_is_set = "--pip-install" in args
+    framework_name2installation = {
+        "numpy": ["numpy"],
+        "torch": ["torch --index-url https://download.pytorch.org/whl/cpu"],
+        "jax": ["jax[cpu]", "flax"],
+        "tensorflow": ["tensorflow"],
+        "cupy": ["cupy"],
+        # switch to stable paddlepaddle, because of https://github.com/PaddlePaddle/Paddle/issues/63927
+        # "paddle": ["paddlepaddle==0.0.0 -f https://www.paddlepaddle.org.cn/whl/linux/cpu-mkl/develop.html"],
+        "paddle": ["paddlepaddle"],
+        "oneflow": ["oneflow==0.9.0"],
+        "pytensor": ["pytensor"],
+    }
+    if sys.platform == "darwin":
+        framework_name2installation["mlx"] = ["mlx"]
+    if sys.platform.startswith("linux"):
+        framework_name2installation["mlx"] = ["mlx[cpu]"]
+
+    usage = f"""
+    Usage:   python -m einops.tests.run_tests <frameworks> [--pip-install]
+    Example: python -m einops.tests.run_tests numpy pytorch --pip-install
+
+    Available frameworks: {list(framework_name2installation)}
+    When --pip-install is set, auto-installs requirements with pip.
+     (make sure which pip points to right pip)
+    """
+    if len(frameworks) == 0:
+        print(usage)
+        return
+    else:
+        synonyms = {
+            "tf": "tensorflow",
+            "pytorch": "torch",
+            "paddlepaddle": "paddle",
+        }
+        frameworks = [synonyms.get(f, f) for f in frameworks]
+        wrong_frameworks = [f for f in frameworks if f not in framework_name2installation]
+        if wrong_frameworks:
+            print(usage)
+            raise RuntimeError(f"Unrecognized frameworks: {wrong_frameworks}")
+
+    if pip_install_is_set:
+        print("Install testing infra")
+        other_dependencies = ["pytest"]
+        assert 0 == run("pip install {} --progress-bar off -q".format(" ".join(other_dependencies)))
+
+        for framework in frameworks:
+            print(f"Installing {framework}")
+            pip_instructions = framework_name2installation[framework]
+            assert 0 == run("pip install {} --progress-bar off -q".format(" ".join(pip_instructions)))
+
+    # we need to inform testing script which frameworks to use
+    # this is done by setting an envvar EINOPS_TEST_BACKENDS
+    from einops.tests import unparse_backends
+
+    envvar_name, envvar_value = unparse_backends(backend_names=frameworks)
+    return_code = run(
+        "python -m pytest .",
+        **{envvar_name: envvar_value},
+    )
+    assert return_code == 0
+
+
+if __name__ == "__main__":
+    main()

kronos-venv/lib/python3.12/site-packages/einops/tests/test_einsum.py

@@ -0,0 +1,356 @@
+import string
+from typing import Any, Callable
+
+import numpy as np
+import pytest
+
+from einops.einops import EinopsError, _compactify_pattern_for_einsum, einsum
+from einops.tests import collect_test_backends
+
+
+class Arguments:
+    def __init__(self, *args: Any, **kargs: Any):
+        self.args = args
+        self.kwargs = kargs
+
+    def __call__(self, function: Callable):
+        return function(*self.args, **self.kwargs)
+
+
+test_layer_cases = [
+    (
+        Arguments("b c_in h w -> w c_out h b", "c_in c_out", bias_shape=None, c_out=13, c_in=12),
+        (2, 12, 3, 4),
+        (4, 13, 3, 2),
+    ),
+    (
+        Arguments("b c_in h w -> w c_out h b", "c_in c_out", bias_shape="c_out", c_out=13, c_in=12),
+        (2, 12, 3, 4),
+        (4, 13, 3, 2),
+    ),
+    (
+        Arguments("b c_in h w -> w c_in h b", "", bias_shape=None, c_in=12),
+        (2, 12, 3, 4),
+        (4, 12, 3, 2),
+    ),
+    (
+        Arguments("b c_in h w -> b c_out", "c_in h w c_out", bias_shape=None, c_in=12, h=3, w=4, c_out=5),
+        (2, 12, 3, 4),
+        (2, 5),
+    ),
+    (
+        Arguments("b t head c_in -> b t head c_out", "head c_in c_out", bias_shape=None, head=4, c_in=5, c_out=6),
+        (2, 3, 4, 5),
+        (2, 3, 4, 6),
+    ),
+]
+
+
+# Each of the form:
+# (Arguments, true_einsum_pattern, in_shapes, out_shape)
+test_functional_cases = [
+    (
+        # Basic:
+        "b c h w, b w -> b h",
+        "abcd,ad->ac",
+        ((2, 3, 4, 5), (2, 5)),
+        (2, 4),
+    ),
+    (
+        # Three tensors:
+        "b c h w, b w, b c -> b h",
+        "abcd,ad,ab->ac",
+        ((2, 3, 40, 5), (2, 5), (2, 3)),
+        (2, 40),
+    ),
+    (
+        # Ellipsis, and full names:
+        "... one two three, three four five -> ... two five",
+        "...abc,cde->...be",
+        ((32, 5, 2, 3, 4), (4, 5, 6)),
+        (32, 5, 3, 6),
+    ),
+    (
+        # Ellipsis at the end:
+        "one two three ..., three four five -> two five ...",
+        "abc...,cde->be...",
+        ((2, 3, 4, 32, 5), (4, 5, 6)),
+        (3, 6, 32, 5),
+    ),
+    (
+        # Ellipsis on multiple tensors:
+        "... one two three, ... three four five -> ... two five",
+        "...abc,...cde->...be",
+        ((32, 5, 2, 3, 4), (32, 5, 4, 5, 6)),
+        (32, 5, 3, 6),
+    ),
+    (
+        # One tensor, and underscores:
+        "first_tensor second_tensor -> first_tensor",
+        "ab->a",
+        ((5, 4),),
+        (5,),
+    ),
+    (
+        # Trace (repeated index)
+        "i i -> ",
+        "aa->",
+        ((5, 5),),
+        (),
+    ),
+    (
+        # Too many spaces in string:
+        " one  two  ,  three four->two  four  ",
+        "ab,cd->bd",
+        ((2, 3), (4, 5)),
+        (3, 5),
+    ),
+    # The following tests were inspired by numpy's einsum tests
+    # https://github.com/numpy/numpy/blob/v1.23.0/numpy/core/tests/test_einsum.py
+    (
+        # Trace with other indices
+        "i middle i -> middle",
+        "aba->b",
+        ((5, 10, 5),),
+        (10,),
+    ),
+    (
+        # Ellipsis in the middle:
+        "i ... i -> ...",
+        "a...a->...",
+        ((5, 3, 2, 1, 4, 5),),
+        (3, 2, 1, 4),
+    ),
+    (
+        # Product of first and last axes:
+        "i ... i -> i ...",
+        "a...a->a...",
+        ((5, 3, 2, 1, 4, 5),),
+        (5, 3, 2, 1, 4),
+    ),
+    (
+        # Triple diagonal
+        "one one one -> one",
+        "aaa->a",
+        ((5, 5, 5),),
+        (5,),
+    ),
+    (
+        # Axis swap:
+        "i j k -> j i k",
+        "abc->bac",
+        ((1, 2, 3),),
+        (2, 1, 3),
+    ),
+    (
+        # Identity:
+        "... -> ...",
+        "...->...",
+        ((5, 4, 3, 2, 1),),
+        (5, 4, 3, 2, 1),
+    ),
+    (
+        # Elementwise product of three tensors
+        "..., ..., ... -> ...",
+        "...,...,...->...",
+        ((3, 2), (3, 2), (3, 2)),
+        (3, 2),
+    ),
+    (
+        # Basic summation:
+        "index ->",
+        "a->",
+        ((10,)),
+        (()),
+    ),
+]
+
+
+def test_layer():
+    for backend in collect_test_backends(layers=True, symbolic=False):
+        rng = np.random.default_rng()
+        if backend.framework_name in ["tensorflow", "torch", "oneflow", "paddle"]:
+            layer_type = backend.layers().EinMix
+            for args, in_shape, out_shape in test_layer_cases:
+                layer = args(layer_type)
+                print("Running", layer.einsum_pattern, "for", backend.framework_name)
+                input = rng.uniform(size=in_shape).astype("float32")
+                input_framework = backend.from_numpy(input)
+                output_framework = layer(input_framework)
+                output = backend.to_numpy(output_framework)
+                assert output.shape == out_shape
+
+
+valid_backends_functional = [
+    "tensorflow",
+    "torch",
+    "jax",
+    "numpy",
+    "oneflow",
+    "cupy",
+    "tensorflow.keras",
+    "paddle",
+    "pytensor",
+    "mlx",
+]
+
+
+def test_functional():
+    # Functional tests:
+    backends = filter(lambda x: x.framework_name in valid_backends_functional, collect_test_backends())
+    for backend in backends:
+        for einops_pattern, true_pattern, in_shapes, out_shape in test_functional_cases:
+            print(f"Running '{einops_pattern}' for {backend.framework_name}")
+
+            # Create pattern:
+            predicted_pattern = _compactify_pattern_for_einsum(einops_pattern)
+            assert predicted_pattern == true_pattern
+
+            # Generate example data:
+            rstate = np.random.RandomState(0)
+            in_arrays = [rstate.uniform(size=shape).astype("float32") for shape in in_shapes]
+            in_arrays_framework = [backend.from_numpy(array) for array in in_arrays]
+
+            # Loop over whether we call it manually with the backend,
+            # or whether we use `einops.einsum`.
+            for do_manual_call in [True, False]:
+                # Actually run einsum:
+                if do_manual_call:
+                    out_array = backend.einsum(predicted_pattern, *in_arrays_framework)
+                else:
+                    out_array = einsum(*in_arrays_framework, einops_pattern)
+
+                # Check shape:
+                if tuple(out_array.shape) != out_shape:
+                    raise ValueError(f"Expected output shape {out_shape} but got {out_array.shape}")
+
+                # Check values:
+                true_out_array = np.einsum(true_pattern, *in_arrays)
+                predicted_out_array = backend.to_numpy(out_array)
+                np.testing.assert_array_almost_equal(predicted_out_array, true_out_array, decimal=5)
+
+
+def test_functional_symbolic():
+    backends = filter(
+        lambda x: x.framework_name in valid_backends_functional, collect_test_backends(symbolic=True, layers=False)
+    )
+    for backend in backends:
+        for einops_pattern, true_pattern, in_shapes, out_shape in test_functional_cases:
+            print(f"Running '{einops_pattern}' for symbolic {backend.framework_name}")
+            # Create pattern:
+            predicted_pattern = _compactify_pattern_for_einsum(einops_pattern)
+            assert predicted_pattern == true_pattern
+
+            rstate = np.random.RandomState(0)
+            in_syms = [backend.create_symbol(in_shape) for in_shape in in_shapes]
+            in_data = [rstate.uniform(size=in_shape).astype("float32") for in_shape in in_shapes]
+
+            expected_out_data = np.einsum(true_pattern, *in_data)
+
+            for do_manual_call in [True, False]:
+                if do_manual_call:
+                    predicted_out_symbol = backend.einsum(predicted_pattern, *in_syms)
+                else:
+                    predicted_out_symbol = einsum(*in_syms, einops_pattern)
+
+                predicted_out_data = backend.eval_symbol(
+                    predicted_out_symbol,
+                    list(zip(in_syms, in_data)),
+                )
+                if predicted_out_data.shape != out_shape:
+                    raise ValueError(f"Expected output shape {out_shape} but got {predicted_out_data.shape}")
+                np.testing.assert_array_almost_equal(predicted_out_data, expected_out_data, decimal=5)
+
+
+def test_functional_errors():
+    # Specific backend does not matter, as errors are raised
+    # during the pattern creation.
+
+    rstate = np.random.RandomState(0)
+
+    def create_tensor(*shape):
+        return rstate.uniform(size=shape).astype("float32")
+
+    # raise NotImplementedError("Singleton () axes are not yet supported in einsum.")
+    with pytest.raises(NotImplementedError, match="^Singleton"):
+        einsum(
+            create_tensor(5, 1),
+            "i () -> i",
+        )
+
+    # raise NotImplementedError("Shape rearrangement is not yet supported in einsum.")
+    with pytest.raises(NotImplementedError, match="^Shape rearrangement"):
+        einsum(
+            create_tensor(5, 1),
+            "a b -> (a b)",
+        )
+
+    with pytest.raises(NotImplementedError, match="^Shape rearrangement"):
+        einsum(
+            create_tensor(10, 1),
+            "(a b) -> a b",
+        )
+
+    # raise RuntimeError("Encountered empty axis name in einsum.")
+    # raise RuntimeError("Axis name in einsum must be a string.")
+    # ^ Not tested, these are just a failsafe in case an unexpected error occurs.
+
+    # raise NotImplementedError("Anonymous axes are not yet supported in einsum.")
+    with pytest.raises(NotImplementedError, match="^Anonymous axes"):
+        einsum(
+            create_tensor(5, 1),
+            "i 2 -> i",
+        )
+
+    # ParsedExpression error:
+    with pytest.raises(EinopsError, match="^Invalid axis identifier"):
+        einsum(
+            create_tensor(5, 1),
+            "i 2j -> i",
+        )
+
+    # raise ValueError("Einsum pattern must contain '->'.")
+    with pytest.raises(ValueError, match="^Einsum pattern"):
+        einsum(
+            create_tensor(5, 3, 2),
+            "i j k",
+        )
+
+    # raise RuntimeError("Too many axes in einsum.")
+    with pytest.raises(RuntimeError, match="^Too many axes"):
+        einsum(
+            create_tensor(1),
+            " ".join(string.ascii_letters) + " extra ->",
+        )
+
+    # raise RuntimeError("Unknown axis on right side of einsum.")
+    with pytest.raises(RuntimeError, match="^Unknown axis"):
+        einsum(
+            create_tensor(5, 1),
+            "i j -> k",
+        )
+
+    # raise ValueError(
+    # "The last argument passed to `einops.einsum` must be a string,"
+    # " representing the einsum pattern."
+    # )
+    with pytest.raises(ValueError, match="^The last argument"):
+        einsum(
+            "i j k -> i",
+            create_tensor(5, 4, 3),
+        )
+
+    # raise ValueError(
+    #     "`einops.einsum` takes at minimum two arguments: the tensors,"
+    #     " followed by the pattern."
+    # )
+    with pytest.raises(ValueError, match="^`einops.einsum` takes"):
+        einsum(
+            "i j k -> i",
+        )
+    with pytest.raises(ValueError, match="^`einops.einsum` takes"):
+        einsum(
+            create_tensor(5, 1),
+        )
+
+    # TODO: Include check for giving normal einsum pattern rather than einops.

kronos-venv/lib/python3.12/site-packages/einops/tests/test_examples.py

@@ -0,0 +1,297 @@
+import numpy as np
+import pytest
+
+from einops import parse_shape, rearrange, reduce
+from einops.tests import is_backend_tested
+from einops.tests.test_ops import imp_op_backends
+
+
+def test_rearrange_examples():
+    def test1(x):
+        # transpose
+        y = rearrange(x, "b c h w -> b h w c")
+        assert tuple(y.shape) == (10, 30, 40, 20)
+        return y
+
+    def test2(x):
+        # view / reshape
+        y = rearrange(x, "b c h w -> b (c h w)")
+        assert tuple(y.shape) == (10, 20 * 30 * 40)
+        return y
+
+    def test3(x):
+        # depth-to-space
+        y = rearrange(x, "b (c h1 w1) h w -> b c (h h1) (w w1)", h1=2, w1=2)
+        assert tuple(y.shape) == (10, 5, 30 * 2, 40 * 2)
+        return y
+
+    def test4(x):
+        # space-to-depth
+        y = rearrange(x, "b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=2, w1=2)
+        assert tuple(y.shape) == (10, 20 * 4, 30 // 2, 40 // 2)
+        return y
+
+    def test5(x):
+        # simple transposition
+        y = rearrange(x, "b1 sound b2 letter -> b1 b2 sound letter")
+        assert tuple(y.shape) == (10, 30, 20, 40)
+        return y
+
+    def test6(x):
+        # parsing parameters
+        t = rearrange(x, "b c h w -> (b h w) c")
+        t = t[:, ::2]  # replacement for dot-product, just changes size of second axis
+        assert tuple(t.shape) == (10 * 30 * 40, 10)
+
+        y = rearrange(t, "(b h w) c2 -> b c2 h w", **parse_shape(x, "b _ h w"))
+        assert tuple(y.shape) == (10, 10, 30, 40)
+        return y
+
+    def test7(x):
+        # split of embedding into groups
+        y1, y2 = rearrange(x, "b (c g) h w -> g b c h w", g=2)
+        assert tuple(y1.shape) == (10, 10, 30, 40)
+        assert tuple(y2.shape) == (10, 10, 30, 40)
+        return y1 + y2  # only one tensor is expected in output
+
+    def test8(x):
+        # max-pooling
+        y = reduce(x, "b c (h h1) (w w1) -> b c h w", reduction="max", h1=2, w1=2)
+        assert tuple(y.shape) == (10, 20, 30 // 2, 40 // 2)
+        return y
+
+    def test9(x):
+        # squeeze - unsqueeze
+        y = reduce(x, "b c h w -> b c () ()", reduction="max")
+        assert tuple(y.shape) == (10, 20, 1, 1)
+        y = rearrange(y, "b c () () -> c b")
+        assert tuple(y.shape) == (20, 10)
+        return y
+
+    def test10(x):
+        # stack
+        tensors = list(x + 0)  # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185
+        tensors = rearrange(tensors, "b c h w -> b h w c")
+        assert tuple(tensors.shape) == (10, 30, 40, 20)
+        return tensors
+
+    def test11(x):
+        # concatenate
+        tensors = list(x + 0)  # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185
+        tensors = rearrange(tensors, "b c h w -> h (b w) c")
+        assert tuple(tensors.shape) == (30, 10 * 40, 20)
+        return tensors
+
+    def shufflenet(x, convolve, c1, c2):
+        # shufflenet reordering example
+        x = convolve(x)
+        x = rearrange(x, "b (c1 c2) h w-> b (c2 c1) h w", c1=c1, c2=c2)
+        x = convolve(x)
+        return x
+
+    def convolve_strided_1d(x, stride, usual_convolution):
+        x = rearrange(x, "b c t1 t2 -> b c (t1 t2)")  # reduce dimensionality
+        x = rearrange(x, "b c (t stride) -> (stride b) c t", stride=stride)
+        x = usual_convolution(x)
+        x = rearrange(x, "(stride b) c t -> b c (t stride)", stride=stride)
+        return x
+
+    def convolve_strided_2d(x, h_stride, w_stride, usual_convolution):
+        x = rearrange(x, "b c (h hs) (w ws) -> (hs ws b) c h w", hs=h_stride, ws=w_stride)
+        x = usual_convolution(x)
+        x = rearrange(x, "(hs ws b) c h w -> b c (h hs) (w ws)", hs=h_stride, ws=w_stride)
+        return x
+
+    def unet_like_1d(x, usual_convolution):
+        # u-net like steps for increasing / reducing dimensionality
+        x = rearrange(x, "b c t1 t2 -> b c (t1 t2)")  # reduce dimensionality
+        y = rearrange(x, "b c (t dt) -> b (dt c) t", dt=2)
+        y = usual_convolution(y)
+        x = x + rearrange(y, "b (dt c) t -> b c (t dt)", dt=2)
+        return x
+
+    # mock for convolution (works for all backends)
+    def convolve_mock(x):
+        return x
+
+    tests = [
+        test1,
+        test2,
+        test3,
+        test4,
+        test5,
+        test6,
+        test7,
+        test8,
+        test9,
+        test10,
+        test11,
+        lambda x: shufflenet(x, convolve=convolve_mock, c1=4, c2=5),
+        lambda x: convolve_strided_1d(x, stride=2, usual_convolution=convolve_mock),
+        lambda x: convolve_strided_2d(x, h_stride=2, w_stride=2, usual_convolution=convolve_mock),
+        lambda x: unet_like_1d(x, usual_convolution=convolve_mock),
+    ]
+
+    for backend in imp_op_backends:
+        print("testing source_examples for ", backend.framework_name)
+        for test in tests:
+            x = np.arange(10 * 20 * 30 * 40).reshape([10, 20, 30, 40])
+            result1 = test(x)
+            result2 = backend.to_numpy(test(backend.from_numpy(x)))
+            assert np.array_equal(result1, result2)
+
+            # now with strides
+            x = np.arange(10 * 2 * 20 * 3 * 30 * 1 * 40).reshape([10 * 2, 20 * 3, 30 * 1, 40 * 1])
+            # known torch bug - torch doesn't support negative steps
+            last_step = -1 if (backend.framework_name != "torch" and backend.framework_name != "oneflow") else 1
+            indexing_expression = np.index_exp[::2, ::3, ::1, ::last_step]
+            result1 = test(x[indexing_expression])
+            result2 = backend.to_numpy(test(backend.from_numpy(x)[indexing_expression]))
+            assert np.array_equal(result1, result2)
+
+
+def tensor_train_example_numpy():
+    # kept here just for a collection, only tested for numpy
+    # https://arxiv.org/pdf/1509.06569.pdf, (5)
+    x = np.ones([3, 4, 5, 6])
+    rank = 4
+    if np.__version__ < "1.15.0":
+        # numpy.einsum fails here, skip test
+        return
+    # creating appropriate Gs
+    Gs = [np.ones([d, d, rank, rank]) for d in x.shape]
+    Gs[0] = Gs[0][:, :, :1, :]
+    Gs[-1] = Gs[-1][:, :, :, :1]
+
+    # einsum way
+    y = x.reshape((1, *x.shape))
+    for G in Gs:
+        # taking partial results left-to-right
+        # y = numpy.einsum('i j alpha beta, alpha i ...  -> beta ... j', G, y)
+        y = np.einsum("i j a b, a i ...  -> b ... j", G, y)
+    y1 = y.reshape(-1)
+
+    # alternative way
+    y = x.reshape(-1)
+    for G in Gs:
+        i, j, alpha, beta = G.shape
+        y = rearrange(y, "(i rest alpha) -> rest (alpha i)", alpha=alpha, i=i)
+        y = y @ rearrange(G, "i j alpha beta -> (alpha i) (j beta)")
+        y = rearrange(y, "rest (beta j) -> (beta rest j)", beta=beta, j=j)
+    y2 = y
+    assert np.allclose(y1, y2)
+
+    # yet another way
+    y = x
+    for G in Gs:
+        i, j, alpha, beta = G.shape
+        y = rearrange(y, "i ... (j alpha) -> ... j (alpha i)", alpha=alpha, i=i)
+        y = y @ rearrange(G, "i j alpha beta -> (alpha i) (j beta)")
+    y3 = y.reshape(-1)
+    assert np.allclose(y1, y3)
+
+
+def test_pytorch_yolo_fragment():
+    if not is_backend_tested("torch"):
+        pytest.skip()
+
+    import torch
+
+    def old_way(tensor, num_classes, num_anchors, anchors, stride_h, stride_w):
+        # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L28-L44
+        bs = tensor.size(0)
+        in_h = tensor.size(2)
+        in_w = tensor.size(3)
+        scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors]
+
+        prediction = tensor.view(bs, num_anchors, 5 + num_classes, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()
+        # Get outputs
+        x = torch.sigmoid(prediction[..., 0])  # Center x
+        y = torch.sigmoid(prediction[..., 1])  # Center y
+        w = prediction[..., 2]  # Width
+        h = prediction[..., 3]  # Height
+        conf = torch.sigmoid(prediction[..., 4])  # Conf
+        pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.
+
+        # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L70-L92
+        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
+        LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
+        # Calculate offsets for each grid
+        grid_x = (
+            torch.linspace(0, in_w - 1, in_w)
+            .repeat(in_w, 1)
+            .repeat(bs * num_anchors, 1, 1)
+            .view(x.shape)
+            .type(FloatTensor)
+        )
+        grid_y = (
+            torch.linspace(0, in_h - 1, in_h)
+            .repeat(in_h, 1)
+            .t()
+            .repeat(bs * num_anchors, 1, 1)
+            .view(y.shape)
+            .type(FloatTensor)
+        )
+        # Calculate anchor w, h
+        anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
+        anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
+        anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape)
+        anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape)
+        # Add offset and scale with anchors
+        pred_boxes = FloatTensor(prediction[..., :4].shape)
+        pred_boxes[..., 0] = x.data + grid_x
+        pred_boxes[..., 1] = y.data + grid_y
+        pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
+        pred_boxes[..., 3] = torch.exp(h.data) * anchor_h
+        # Results
+        _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)
+        output = torch.cat(
+            (pred_boxes.view(bs, -1, 4) * _scale, conf.view(bs, -1, 1), pred_cls.view(bs, -1, num_classes)), -1
+        )
+        return output
+
+    def new_way(tensor, num_classes, num_anchors, anchors, stride_h, stride_w):
+        raw_predictions = rearrange(tensor, " b (anchor prediction) h w -> prediction b anchor h w", anchor=num_anchors)
+
+        anchors = torch.FloatTensor(anchors).to(tensor.device)
+        anchor_sizes = rearrange(anchors, "anchor dim -> dim () anchor () ()")
+
+        _, _, _, in_h, in_w = raw_predictions.shape
+        grid_h = rearrange(torch.arange(in_h).float(), "h -> () () h ()").to(tensor.device)
+        grid_w = rearrange(torch.arange(in_w).float(), "w -> () () () w").to(tensor.device)
+
+        predicted_bboxes = torch.zeros_like(raw_predictions)
+        predicted_bboxes[0] = (raw_predictions[0].sigmoid() + grid_h) * stride_h  # center y
+        predicted_bboxes[1] = (raw_predictions[1].sigmoid() + grid_w) * stride_w  # center x
+        predicted_bboxes[2:4] = (raw_predictions[2:4].exp()) * anchor_sizes  # bbox width and height
+        predicted_bboxes[4] = raw_predictions[4].sigmoid()  # confidence
+        predicted_bboxes[5:] = raw_predictions[5:].sigmoid()  # class predictions
+        # only to match results of original code, not needed
+        return rearrange(predicted_bboxes, "prediction b anchor h w -> b anchor h w prediction")
+
+    stride_h = 4
+    stride_w = 4
+    batch_size = 5
+    num_classes = 12
+    anchors = [[50, 100], [100, 50], [75, 75]]
+    num_anchors = len(anchors)
+
+    x = torch.randn([batch_size, num_anchors * (5 + num_classes), 1, 1])
+    result1 = old_way(
+        tensor=x,
+        num_anchors=num_anchors,
+        num_classes=num_classes,
+        stride_h=stride_h,
+        stride_w=stride_w,
+        anchors=anchors,
+    )
+    result2 = new_way(
+        tensor=x,
+        num_anchors=num_anchors,
+        num_classes=num_classes,
+        stride_h=stride_h,
+        stride_w=stride_w,
+        anchors=anchors,
+    )
+    result1 = result1.reshape(result2.shape)
+    assert torch.allclose(result1, result2)

kronos-venv/lib/python3.12/site-packages/einops/tests/test_layers.py

@@ -0,0 +1,480 @@
+import pickle
+from collections import namedtuple
+
+import numpy as np
+import pytest
+
+from einops import EinopsError, rearrange, reduce
+from einops.tests import FLOAT_REDUCTIONS as REDUCTIONS
+from einops.tests import collect_test_backends, is_backend_tested
+
+__author__ = "Alex Rogozhnikov"
+
+testcase = namedtuple("testcase", ["pattern", "axes_lengths", "input_shape", "wrong_shapes"])
+
+rearrangement_patterns = [
+    testcase(
+        "b c h w -> b (c h w)",
+        dict(c=20),
+        (10, 20, 30, 40),
+        [(), (10,), (10, 10, 10), (10, 21, 30, 40), [1, 20, 1, 1, 1]],
+    ),
+    testcase(
+        "b c (h1 h2) (w1 w2) -> b (c h2 w2) h1 w1",
+        dict(h2=2, w2=2),
+        (10, 20, 30, 40),
+        [(), (1, 1, 1, 1), (1, 10, 3), ()],
+    ),
+    testcase(
+        "b ... c -> c b ...",
+        dict(b=10),
+        (10, 20, 30),
+        [(), (10,), (5, 10)],
+    ),
+]
+
+
+def test_rearrange_imperative():
+    for backend in collect_test_backends(symbolic=False, layers=True):
+        print("Test layer for ", backend.framework_name)
+
+        for pattern, axes_lengths, input_shape, wrong_shapes in rearrangement_patterns:
+            x = np.arange(np.prod(input_shape), dtype="float32").reshape(input_shape)
+            result_numpy = rearrange(x, pattern, **axes_lengths)
+            layer = backend.layers().Rearrange(pattern, **axes_lengths)
+            for shape in wrong_shapes:
+                try:
+                    layer(backend.from_numpy(np.zeros(shape, dtype="float32")))
+                except BaseException:
+                    pass
+                else:
+                    raise AssertionError("Failure expected")
+
+            # simple pickling / unpickling
+            layer2 = pickle.loads(pickle.dumps(layer))
+            result1 = backend.to_numpy(layer(backend.from_numpy(x)))
+            result2 = backend.to_numpy(layer2(backend.from_numpy(x)))
+            assert np.allclose(result_numpy, result1)
+            assert np.allclose(result1, result2)
+
+            just_sum = backend.layers().Reduce("...->", reduction="sum")
+
+            variable = backend.from_numpy(x)
+            result = just_sum(layer(variable))
+
+            result.backward()
+            assert np.allclose(backend.to_numpy(variable.grad), 1)
+
+
+def test_rearrange_symbolic():
+    for backend in collect_test_backends(symbolic=True, layers=True):
+        print("Test layer for ", backend.framework_name)
+
+        for pattern, axes_lengths, input_shape, _wrong_shapes in rearrangement_patterns:
+            x = np.arange(np.prod(input_shape), dtype="float32").reshape(input_shape)
+            result_numpy = rearrange(x, pattern, **axes_lengths)
+            layer = backend.layers().Rearrange(pattern, **axes_lengths)
+            input_shape_of_nones = [None] * len(input_shape)
+            shapes = [input_shape, input_shape_of_nones]
+
+            for shape in shapes:
+                symbol = backend.create_symbol(shape)
+                eval_inputs = [(symbol, x)]
+
+                result_symbol1 = layer(symbol)
+                result1 = backend.eval_symbol(result_symbol1, eval_inputs)
+                assert np.allclose(result_numpy, result1)
+
+                layer2 = pickle.loads(pickle.dumps(layer))
+                result_symbol2 = layer2(symbol)
+                result2 = backend.eval_symbol(result_symbol2, eval_inputs)
+                assert np.allclose(result1, result2)
+
+                # now testing back-propagation
+                just_sum = backend.layers().Reduce("...->", reduction="sum")
+
+                result_sum1 = backend.eval_symbol(just_sum(result_symbol1), eval_inputs)
+                result_sum2 = np.sum(x)
+
+                assert np.allclose(result_sum1, result_sum2)
+
+
+reduction_patterns = [
+    *rearrangement_patterns,
+    testcase("b c h w -> b ()", dict(b=10), (10, 20, 30, 40), [(10,), (10, 20, 30)]),
+    testcase("b c (h1 h2) (w1 w2) -> b c h1 w1", dict(h1=15, h2=2, w2=2), (10, 20, 30, 40), [(10, 20, 31, 40)]),
+    testcase("b ... c -> b", dict(b=10), (10, 20, 30, 40), [(10,), (11, 10)]),
+]
+
+
+def test_reduce_imperative():
+    for backend in collect_test_backends(symbolic=False, layers=True):
+        print("Test layer for ", backend.framework_name)
+        for reduction in REDUCTIONS:
+            for pattern, axes_lengths, input_shape, wrong_shapes in reduction_patterns:
+                print(backend, reduction, pattern, axes_lengths, input_shape, wrong_shapes)
+                x = np.arange(1, 1 + np.prod(input_shape), dtype="float32").reshape(input_shape)
+                x /= x.mean()
+                result_numpy = reduce(x, pattern, reduction, **axes_lengths)
+                layer = backend.layers().Reduce(pattern, reduction, **axes_lengths)
+                for shape in wrong_shapes:
+                    try:
+                        layer(backend.from_numpy(np.zeros(shape, dtype="float32")))
+                    except BaseException:
+                        pass
+                    else:
+                        raise AssertionError("Failure expected")
+
+                # simple pickling / unpickling
+                layer2 = pickle.loads(pickle.dumps(layer))
+                result1 = backend.to_numpy(layer(backend.from_numpy(x)))
+                result2 = backend.to_numpy(layer2(backend.from_numpy(x)))
+                assert np.allclose(result_numpy, result1)
+                assert np.allclose(result1, result2)
+
+                just_sum = backend.layers().Reduce("...->", reduction="sum")
+
+                variable = backend.from_numpy(x)
+                result = just_sum(layer(variable))
+
+                result.backward()
+                grad = backend.to_numpy(variable.grad)
+                if reduction == "sum":
+                    assert np.allclose(grad, 1)
+                if reduction == "mean":
+                    assert np.allclose(grad, grad.min())
+                if reduction in ["max", "min"]:
+                    assert np.all(np.isin(grad, [0, 1]))
+                    assert np.sum(grad) > 0.5
+
+
+def test_reduce_symbolic():
+    for backend in collect_test_backends(symbolic=True, layers=True):
+        print("Test layer for ", backend.framework_name)
+        for reduction in REDUCTIONS:
+            for pattern, axes_lengths, input_shape, _wrong_shapes in reduction_patterns:
+                x = np.arange(1, 1 + np.prod(input_shape), dtype="float32").reshape(input_shape)
+                x /= x.mean()
+                result_numpy = reduce(x, pattern, reduction, **axes_lengths)
+                layer = backend.layers().Reduce(pattern, reduction, **axes_lengths)
+                input_shape_of_nones = [None] * len(input_shape)
+                shapes = [input_shape, input_shape_of_nones]
+
+                for shape in shapes:
+                    symbol = backend.create_symbol(shape)
+                    eval_inputs = [(symbol, x)]
+
+                    result_symbol1 = layer(symbol)
+                    result1 = backend.eval_symbol(result_symbol1, eval_inputs)
+                    assert np.allclose(result_numpy, result1)
+
+                    layer2 = pickle.loads(pickle.dumps(layer))
+                    result_symbol2 = layer2(symbol)
+                    result2 = backend.eval_symbol(result_symbol2, eval_inputs)
+                    assert np.allclose(result1, result2)
+
+
+def create_torch_model(use_reduce=False, add_scripted_layer=False):
+    if not is_backend_tested("torch"):
+        pytest.skip()
+    else:
+        import torch.jit
+        from torch.nn import Conv2d, Linear, MaxPool2d, ReLU, Sequential
+
+        from einops.layers.torch import EinMix, Rearrange, Reduce
+
+        return Sequential(
+            Conv2d(3, 6, kernel_size=(5, 5)),
+            Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2) if use_reduce else MaxPool2d(kernel_size=2),
+            Conv2d(6, 16, kernel_size=(5, 5)),
+            Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2),
+            torch.jit.script(Rearrange("b c h w -> b (c h w)"))
+            if add_scripted_layer
+            else Rearrange("b c h w -> b (c h w)"),
+            Linear(16 * 5 * 5, 120),
+            ReLU(),
+            Linear(120, 84),
+            ReLU(),
+            EinMix("b c1 -> (b c2)", weight_shape="c1 c2", bias_shape="c2", c1=84, c2=84),
+            EinMix("(b c2) -> b c3", weight_shape="c2 c3", bias_shape="c3", c2=84, c3=84),
+            Linear(84, 10),
+        )
+
+
+def test_torch_layer():
+    if not is_backend_tested("torch"):
+        pytest.skip()
+    else:
+        # checked that torch present
+        import torch
+        import torch.jit
+
+        model1 = create_torch_model(use_reduce=True)
+        model2 = create_torch_model(use_reduce=False)
+        input = torch.randn([10, 3, 32, 32])
+        # random models have different predictions
+        assert not torch.allclose(model1(input), model2(input))
+        model2.load_state_dict(pickle.loads(pickle.dumps(model1.state_dict())))
+        assert torch.allclose(model1(input), model2(input))
+
+        # tracing (freezing)
+        model3 = torch.jit.trace(model2, example_inputs=input)
+        torch.testing.assert_close(model1(input), model3(input), atol=1e-3, rtol=1e-3)
+        torch.testing.assert_close(model1(input + 1), model3(input + 1), atol=1e-3, rtol=1e-3)
+
+        model4 = torch.jit.trace(model2, example_inputs=input)
+        torch.testing.assert_close(model1(input), model4(input), atol=1e-3, rtol=1e-3)
+        torch.testing.assert_close(model1(input + 1), model4(input + 1), atol=1e-3, rtol=1e-3)
+
+
+def test_torch_layers_scripting():
+    if not is_backend_tested("torch"):
+        pytest.skip()
+    else:
+        import torch
+
+        for script_layer in [False, True]:
+            model1 = create_torch_model(use_reduce=True, add_scripted_layer=script_layer)
+            model2 = torch.jit.script(model1)
+            input = torch.randn([10, 3, 32, 32])
+
+            torch.testing.assert_close(model1(input), model2(input), atol=1e-3, rtol=1e-3)
+
+
+def test_keras_layer():
+    rng = np.random.default_rng()
+    if not is_backend_tested("tensorflow"):
+        pytest.skip()
+    else:
+        import tensorflow as tf
+
+        if tf.__version__ < "2.16.":
+            # current implementation of layers follows new TF interface
+            pytest.skip()
+        from tensorflow.keras.layers import Conv2D as Conv2d
+        from tensorflow.keras.layers import Dense as Linear
+        from tensorflow.keras.layers import ReLU
+        from tensorflow.keras.models import Sequential
+
+        from einops.layers.keras import EinMix, Rearrange, Reduce, keras_custom_objects
+
+        def create_keras_model():
+            return Sequential(
+                [
+                    Conv2d(6, kernel_size=5, input_shape=[32, 32, 3]),
+                    Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2),
+                    Conv2d(16, kernel_size=5),
+                    Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2),
+                    Rearrange("b c h w -> b (c h w)"),
+                    Linear(120),
+                    ReLU(),
+                    Linear(84),
+                    ReLU(),
+                    EinMix("b c1 -> (b c2)", weight_shape="c1 c2", bias_shape="c2", c1=84, c2=84),
+                    EinMix("(b c2) -> b c3", weight_shape="c2 c3", bias_shape="c3", c2=84, c3=84),
+                    Linear(10),
+                ]
+            )
+
+        model1 = create_keras_model()
+        model2 = create_keras_model()
+
+        input = rng.normal(size=[10, 32, 32, 3]).astype("float32")
+        # two randomly init models should provide different outputs
+        assert not np.allclose(model1.predict_on_batch(input), model2.predict_on_batch(input))
+
+        # get some temp filename
+        tmp_model_filename = "/tmp/einops_tf_model.h5"
+        # save arch + weights
+        print("temp_path_keras1", tmp_model_filename)
+        tf.keras.models.save_model(model1, tmp_model_filename)
+        model3 = tf.keras.models.load_model(tmp_model_filename, custom_objects=keras_custom_objects)
+
+        np.testing.assert_allclose(model1.predict_on_batch(input), model3.predict_on_batch(input))
+
+        weight_filename = "/tmp/einops_tf_model.weights.h5"
+        # save arch as json
+        model4 = tf.keras.models.model_from_json(model1.to_json(), custom_objects=keras_custom_objects)
+        model1.save_weights(weight_filename)
+        model4.load_weights(weight_filename)
+        model2.load_weights(weight_filename)
+        # check that differently-inialized model receives same weights
+        np.testing.assert_allclose(model1.predict_on_batch(input), model2.predict_on_batch(input))
+        # ulimate test
+        # save-load architecture, and then load weights - should return same result
+        np.testing.assert_allclose(model1.predict_on_batch(input), model4.predict_on_batch(input))
+
+
+def test_flax_layers():
+    """
+    One-off simple tests for Flax layers.
+    Unfortunately, Flax layers have a different interface from other layers.
+    """
+    if not is_backend_tested("jax"):
+        pytest.skip()
+    else:
+        import flax
+        import jax
+        import jax.numpy as jnp
+        from flax import linen as nn
+
+        from einops.layers.flax import EinMix, Rearrange, Reduce
+
+        class NN(nn.Module):
+            @nn.compact
+            def __call__(self, x):
+                x = EinMix(
+                    "b (h h2) (w w2) c -> b h w c_out", "h2 w2 c c_out", "c_out", sizes=dict(h2=2, w2=3, c=4, c_out=5)
+                )(x)
+                x = Rearrange("b h w c -> b (w h c)", sizes=dict(c=5))(x)
+                x = Reduce("b hwc -> b", "mean", dict(hwc=2 * 3 * 5))(x)
+                return x
+
+        model = NN()
+        fixed_input = jnp.ones([10, 2 * 2, 3 * 3, 4])
+        params = model.init(jax.random.PRNGKey(0), fixed_input)
+
+        def eval_at_point(params):
+            return jnp.linalg.norm(model.apply(params, fixed_input))
+
+        vandg = jax.value_and_grad(eval_at_point)
+        value0 = eval_at_point(params)
+        value1, grad1 = vandg(params)
+        assert jnp.allclose(value0, value1)
+        if jax.__version__ < "0.6.0":
+            tree_map = jax.tree_map
+        else:
+            tree_map = jax.tree.map
+
+        params2 = tree_map(lambda x1, x2: x1 - x2 * 0.001, params, grad1)
+
+        value2 = eval_at_point(params2)
+        assert value0 >= value2, (value0, value2)
+
+        # check serialization
+        fbytes = flax.serialization.to_bytes(params)
+        _loaded = flax.serialization.from_bytes(params, fbytes)
+
+
+def test_einmix_decomposition():
+    """
+    Testing that einmix correctly decomposes into smaller transformations.
+    """
+    from einops.layers._einmix import _EinmixDebugger
+
+    mixin1 = _EinmixDebugger(
+        "a b c d e -> e d c b a",
+        weight_shape="d a b",
+        d=2, a=3, b=5,
+    )  # fmt: off
+    assert mixin1.pre_reshape_pattern is None
+    assert mixin1.post_reshape_pattern is None
+    assert mixin1.einsum_pattern == "abcde,dab->edcba"
+    assert mixin1.saved_weight_shape == [2, 3, 5]
+    assert mixin1.saved_bias_shape is None
+
+    mixin2 = _EinmixDebugger(
+        "a b c d e -> e d c b a",
+        weight_shape="d a b",
+        bias_shape="a b c d e",
+        a=1, b=2, c=3, d=4, e=5,
+    )  # fmt: off
+    assert mixin2.pre_reshape_pattern is None
+    assert mixin2.post_reshape_pattern is None
+    assert mixin2.einsum_pattern == "abcde,dab->edcba"
+    assert mixin2.saved_weight_shape == [4, 1, 2]
+    assert mixin2.saved_bias_shape == [5, 4, 3, 2, 1]
+
+    mixin3 = _EinmixDebugger(
+        "... -> ...",
+        weight_shape="",
+        bias_shape="",
+    )  # fmt: off
+    assert mixin3.pre_reshape_pattern is None
+    assert mixin3.post_reshape_pattern is None
+    assert mixin3.einsum_pattern == "...,->..."
+    assert mixin3.saved_weight_shape == []
+    assert mixin3.saved_bias_shape == []
+
+    mixin4 = _EinmixDebugger(
+        "b a ...  -> b c ...",
+        weight_shape="b a c",
+        a=1, b=2, c=3,
+    )  # fmt: off
+    assert mixin4.pre_reshape_pattern is None
+    assert mixin4.post_reshape_pattern is None
+    assert mixin4.einsum_pattern == "ba...,bac->bc..."
+    assert mixin4.saved_weight_shape == [2, 1, 3]
+    assert mixin4.saved_bias_shape is None
+
+    mixin5 = _EinmixDebugger(
+        "(b a) ... -> b c (...)",
+        weight_shape="b a c",
+        a=1, b=2, c=3,
+    )  # fmt: off
+    assert mixin5.pre_reshape_pattern == "(b a) ... -> b a ..."
+    assert mixin5.pre_reshape_lengths == dict(a=1, b=2)
+    assert mixin5.post_reshape_pattern == "b c ... -> b c (...)"
+    assert mixin5.einsum_pattern == "ba...,bac->bc..."
+    assert mixin5.saved_weight_shape == [2, 1, 3]
+    assert mixin5.saved_bias_shape is None
+
+    mixin6 = _EinmixDebugger(
+        "b ... (a c) -> b ... (a d)",
+        weight_shape="c d",
+        bias_shape="a d",
+        a=1, c=3, d=4,
+    )  # fmt: off
+    assert mixin6.pre_reshape_pattern == "b ... (a c) -> b ... a c"
+    assert mixin6.pre_reshape_lengths == dict(a=1, c=3)
+    assert mixin6.post_reshape_pattern == "b ... a d -> b ... (a d)"
+    assert mixin6.einsum_pattern == "b...ac,cd->b...ad"
+    assert mixin6.saved_weight_shape == [3, 4]
+    assert mixin6.saved_bias_shape == [1, 1, 4]  # (b) a d, ellipsis does not participate
+
+    mixin7 = _EinmixDebugger(
+        "a ... (b c) -> a (... d b)",
+        weight_shape="c d b",
+        bias_shape="d b",
+        b=2, c=3, d=4,
+    )  # fmt: off
+    assert mixin7.pre_reshape_pattern == "a ... (b c) -> a ... b c"
+    assert mixin7.pre_reshape_lengths == dict(b=2, c=3)
+    assert mixin7.post_reshape_pattern == "a ... d b -> a (... d b)"
+    assert mixin7.einsum_pattern == "a...bc,cdb->a...db"
+    assert mixin7.saved_weight_shape == [3, 4, 2]
+    assert mixin7.saved_bias_shape == [1, 4, 2]  # (a) d b, ellipsis does not participate
+
+
+def test_einmix_restrictions():
+    """
+    Testing different cases
+    """
+    from einops.layers._einmix import _EinmixDebugger
+
+    with pytest.raises(EinopsError):
+        _EinmixDebugger(
+            "a b c d e -> e d c b a",
+            weight_shape="d a b",
+            d=2, a=3, # missing b
+        )  # fmt: off
+
+    with pytest.raises(EinopsError):
+        _EinmixDebugger(
+            "a b c d e -> e d c b a",
+            weight_shape="w a b",
+            d=2, a=3, b=1 # missing d
+        )  # fmt: off
+
+    with pytest.raises(EinopsError):
+        _EinmixDebugger(
+            "(...) a -> ... a",
+            weight_shape="a", a=1, # ellipsis on the left
+        )  # fmt: off
+
+    with pytest.raises(EinopsError):
+        _EinmixDebugger(
+            "(...) a -> a ...",
+            weight_shape="a", a=1, # ellipsis on the right side after bias axis
+            bias_shape="a",
+        )  # fmt: off

kronos-venv/lib/python3.12/site-packages/einops/tests/test_ops.py

@@ -0,0 +1,658 @@
+import itertools
+
+import numpy as np
+import pytest
+
+from einops import EinopsError
+from einops.einops import _enumerate_directions, rearrange, reduce, repeat
+from einops.tests import FLOAT_REDUCTIONS as REDUCTIONS
+from einops.tests import collect_test_backends, is_backend_tested
+
+imp_op_backends = collect_test_backends(symbolic=False, layers=False)
+sym_op_backends = collect_test_backends(symbolic=True, layers=False)
+
+rng = np.random.default_rng()
+
+identity_patterns = [
+    "...->...",
+    "a b c d e-> a b c d e",
+    "a b c d e ...-> ... a b c d e",
+    "a b c d e ...-> a ... b c d e",
+    "... a b c d e -> ... a b c d e",
+    "a ... e-> a ... e",
+    "a ... -> a ... ",
+    "a ... c d e -> a (...) c d e",
+]
+
+equivalent_rearrange_patterns = [
+    ("a b c d e -> (a b) c d e", "a b ... -> (a b) ... "),
+    ("a b c d e -> a b (c d) e", "... c d e -> ... (c d) e"),
+    ("a b c d e -> a b c d e", "... -> ... "),
+    ("a b c d e -> (a b c d e)", "... ->  (...)"),
+    ("a b c d e -> b (c d e) a", "a b ... -> b (...) a"),
+    ("a b c d e -> b (a c d) e", "a b ... e -> b (a ...) e"),
+]
+
+equivalent_reduction_patterns = [
+    ("a b c d e -> ", " ... ->  "),
+    ("a b c d e -> (e a)", "a ... e -> (e a)"),
+    ("a b c d e -> d (a e)", " a b c d e ... -> d (a e) "),
+    ("a b c d e -> (a b)", " ... c d e  -> (...) "),
+]
+
+
+def test_collapsed_ellipsis_errors_out():
+    x = np.zeros([1, 1, 1, 1, 1])
+    rearrange(x, "a b c d ... ->  a b c ... d")
+    with pytest.raises(EinopsError):
+        rearrange(x, "a b c d (...) ->  a b c ... d")
+
+    rearrange(x, "... ->  (...)")
+    with pytest.raises(EinopsError):
+        rearrange(x, "(...) -> (...)")
+
+
+def test_ellipsis_ops_numpy():
+    x = np.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
+    for pattern in identity_patterns:
+        assert np.array_equal(x, rearrange(x, pattern)), pattern
+
+    for pattern1, pattern2 in equivalent_rearrange_patterns:
+        assert np.array_equal(rearrange(x, pattern1), rearrange(x, pattern2))
+
+    for reduction in ["min", "max", "sum"]:
+        for pattern1, pattern2 in equivalent_reduction_patterns:
+            assert np.array_equal(reduce(x, pattern1, reduction=reduction), reduce(x, pattern2, reduction=reduction))
+
+    # now just check coincidence with numpy
+    all_rearrange_patterns = [*identity_patterns]
+    for pattern_pairs in equivalent_rearrange_patterns:
+        all_rearrange_patterns.extend(pattern_pairs)
+
+
+def check_op_against_numpy(backend, numpy_input, pattern, axes_lengths, reduction="rearrange", is_symbolic=False):
+    """
+    Helper to test result of operation (rearrange or transpose) against numpy
+    if reduction == 'rearrange', rearrange op is tested, otherwise reduce
+    """
+
+    def operation(x):
+        if reduction == "rearrange":
+            return rearrange(x, pattern, **axes_lengths)
+        else:
+            return reduce(x, pattern, reduction, **axes_lengths)
+
+    numpy_result = operation(numpy_input)
+    check_equal = np.array_equal
+    p_none_dimension = 0.5
+    if is_symbolic:
+        symbol_shape = [d if rng.random() >= p_none_dimension else None for d in numpy_input.shape]
+        symbol = backend.create_symbol(shape=symbol_shape)
+        result_symbol = operation(symbol)
+        backend_result = backend.eval_symbol(result_symbol, [(symbol, numpy_input)])
+    else:
+        backend_result = operation(backend.from_numpy(numpy_input))
+        backend_result = backend.to_numpy(backend_result)
+
+    check_equal(numpy_result, backend_result)
+
+
+def test_ellipsis_ops_imperative():
+    """Checking various patterns against numpy"""
+    x = np.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
+    for is_symbolic in [True, False]:
+        for backend in collect_test_backends(symbolic=is_symbolic, layers=False):
+            for pattern in identity_patterns + list(itertools.chain(*equivalent_rearrange_patterns)):
+                check_op_against_numpy(
+                    backend, x, pattern, axes_lengths={}, reduction="rearrange", is_symbolic=is_symbolic
+                )
+
+            for reduction in ["min", "max", "sum"]:
+                for pattern in itertools.chain(*equivalent_reduction_patterns):
+                    check_op_against_numpy(
+                        backend, x, pattern, axes_lengths={}, reduction=reduction, is_symbolic=is_symbolic
+                    )
+
+
+def test_rearrange_array_api():
+    import numpy as xp
+
+    from einops import array_api as AA
+
+    if xp.__version__ < "2.0.0":
+        pytest.skip()
+
+    x = np.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
+    for pattern in identity_patterns + list(itertools.chain(*equivalent_rearrange_patterns)):
+        expected = rearrange(x, pattern)
+        result = AA.rearrange(xp.from_dlpack(x), pattern)
+        assert np.array_equal(AA.asnumpy(result + 0), expected)
+
+
+def test_reduce_array_api():
+    import numpy as xp
+
+    from einops import array_api as AA
+
+    if xp.__version__ < "2.0.0":
+        pytest.skip()
+
+    x = np.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
+    for pattern in itertools.chain(*equivalent_reduction_patterns):
+        for reduction in ["min", "max", "sum"]:
+            expected = reduce(x, pattern, reduction=reduction)
+            result = AA.reduce(xp.from_dlpack(x), pattern, reduction=reduction)
+            assert np.array_equal(AA.asnumpy(np.asarray(result + 0)), expected)
+
+
+def test_rearrange_consistency_numpy():
+    shape = [1, 2, 3, 5, 7, 11]
+    x = np.arange(np.prod(shape)).reshape(shape)
+    for pattern in [
+        "a b c d e f -> a b c d e f",
+        "b a c d e f -> a b d e f c",
+        "a b c d e f -> f e d c b a",
+        "a b c d e f -> (f e) d (c b a)",
+        "a b c d e f -> (f e d c b a)",
+    ]:
+        result = rearrange(x, pattern)
+        assert len(np.setdiff1d(x, result)) == 0
+        assert result.dtype == x.dtype
+
+    result = rearrange(x, "a b c d e f -> a (b) (c d e) f")
+    assert np.array_equal(x.flatten(), result.flatten())
+
+    result = rearrange(x, "a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11")
+    assert np.array_equal(x, result)
+
+    result1 = rearrange(x, "a b c d e f -> f e d c b a")
+    result2 = rearrange(x, "f e d c b a -> a b c d e f")
+    assert np.array_equal(result1, result2)
+
+    result = rearrange(rearrange(x, "a b c d e f -> (f d) c (e b) a"), "(f d) c (e b) a -> a b c d e f", b=2, d=5)
+    assert np.array_equal(x, result)
+
+    sizes = dict(zip("abcdef", shape))
+    temp = rearrange(x, "a b c d e f -> (f d) c (e b) a", **sizes)
+    result = rearrange(temp, "(f d) c (e b) a -> a b c d e f", **sizes)
+    assert np.array_equal(x, result)
+
+    x2 = np.arange(2 * 3 * 4).reshape([2, 3, 4])
+    result = rearrange(x2, "a b c -> b c a")
+    assert x2[1, 2, 3] == result[2, 3, 1]
+    assert x2[0, 1, 2] == result[1, 2, 0]
+
+
+def test_rearrange_permutations_numpy():
+    # tests random permutation of axes against two independent numpy ways
+    for n_axes in range(1, 10):
+        input = np.arange(2**n_axes).reshape([2] * n_axes)
+        permutation = rng.permutation(n_axes)
+        left_expression = " ".join("i" + str(axis) for axis in range(n_axes))
+        right_expression = " ".join("i" + str(axis) for axis in permutation)
+        expression = left_expression + " -> " + right_expression
+        result = rearrange(input, expression)
+
+        for pick in rng.integers(0, 2, [10, n_axes]):
+            assert input[tuple(pick)] == result[tuple(pick[permutation])]
+
+    for n_axes in range(1, 10):
+        input = np.arange(2**n_axes).reshape([2] * n_axes)
+        permutation = rng.permutation(n_axes)
+        left_expression = " ".join("i" + str(axis) for axis in range(n_axes)[::-1])
+        right_expression = " ".join("i" + str(axis) for axis in permutation[::-1])
+        expression = left_expression + " -> " + right_expression
+        result = rearrange(input, expression)
+        assert result.shape == input.shape
+        expected_result = np.zeros_like(input)
+        for original_axis, result_axis in enumerate(permutation):
+            expected_result |= ((input >> original_axis) & 1) << result_axis
+
+        assert np.array_equal(result, expected_result)
+
+
+def test_reduction_imperatives():
+    for backend in imp_op_backends:
+        print("Reduction tests for ", backend.framework_name)
+        for reduction in REDUCTIONS:
+            # slight redundancy for simpler order - numpy version is evaluated multiple times
+            input = np.arange(2 * 3 * 4 * 5 * 6, dtype="int64").reshape([2, 3, 4, 5, 6])
+            if reduction in ["mean", "prod"]:
+                input = input / input.astype("float64").mean()
+            test_cases = [
+                ["a b c d e -> ", {}, getattr(input, reduction)()],
+                ["a ... -> ", {}, getattr(input, reduction)()],
+                ["(a1 a2) ... (e1 e2) -> ", dict(a1=1, e2=2), getattr(input, reduction)()],
+                [
+                    "a b c d e -> (e c) a",
+                    {},
+                    getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]),
+                ],
+                [
+                    "a ... c d e -> (e c) a",
+                    {},
+                    getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]),
+                ],
+                [
+                    "a b c d e ... -> (e c) a",
+                    {},
+                    getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]),
+                ],
+                ["a b c d e -> (e c a)", {}, getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1])],
+                ["(a a2) ... -> (a2 a) ...", dict(a2=1), input],
+            ]
+            for pattern, axes_lengths, expected_result in test_cases:
+                result = reduce(backend.from_numpy(input.copy()), pattern, reduction=reduction, **axes_lengths)
+                result = backend.to_numpy(result)
+                assert np.allclose(result, expected_result), f"Failed at {pattern}"
+
+
+def test_reduction_symbolic():
+    for backend in sym_op_backends:
+        print("Reduction tests for ", backend.framework_name)
+        for reduction in REDUCTIONS:
+            input = np.arange(2 * 3 * 4 * 5 * 6, dtype="int64").reshape([2, 3, 4, 5, 6])
+            input = input / input.astype("float64").mean()
+            # slight redundancy for simpler order - numpy version is evaluated multiple times
+            test_cases = [
+                ["a b c d e -> ", {}, getattr(input, reduction)()],
+                ["a ... -> ", {}, getattr(input, reduction)()],
+                ["(a a2) ... (e e2) -> ", dict(a2=1, e2=1), getattr(input, reduction)()],
+                [
+                    "a b c d e -> (e c) a",
+                    {},
+                    getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]),
+                ],
+                [
+                    "a ... c d e -> (e c) a",
+                    {},
+                    getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]),
+                ],
+                [
+                    "a b c d e ... -> (e c) a",
+                    {},
+                    getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]),
+                ],
+                ["a b c d e -> (e c a)", {}, getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1])],
+                ["(a a2) ... -> (a2 a) ...", dict(a2=1), input],
+            ]
+            for pattern, axes_lengths, expected_numpy_result in test_cases:
+                shapes = [input.shape, [None for _ in input.shape]]
+                for shape in shapes:
+                    sym = backend.create_symbol(shape)
+                    result_sym = reduce(sym, pattern, reduction=reduction, **axes_lengths)
+                    result = backend.eval_symbol(result_sym, [(sym, input)])
+                    assert np.allclose(result, expected_numpy_result)
+
+                if True:
+                    shape = []
+                    _axes_lengths = {**axes_lengths}
+                    for axis, length in zip("abcde", input.shape):
+                        # filling as much as possible with Nones
+                        if axis in pattern:
+                            shape.append(None)
+                            _axes_lengths[axis] = length
+                        else:
+                            shape.append(length)
+                    sym = backend.create_symbol(shape)
+                    result_sym = reduce(sym, pattern, reduction=reduction, **_axes_lengths)
+                    result = backend.eval_symbol(result_sym, [(sym, input)])
+                    assert np.allclose(result, expected_numpy_result)
+
+
+def test_reduction_stress_imperatives():
+    for backend in imp_op_backends:
+        print("Stress-testing reduction for ", backend.framework_name)
+        for reduction in [*REDUCTIONS, "rearrange"]:
+            dtype = "int64"
+            coincide = np.array_equal
+            if reduction in ["mean", "prod"]:
+                dtype = "float64"
+                coincide = np.allclose
+            max_dim = 11
+            if "oneflow" in backend.framework_name:
+                max_dim = 7
+            if "paddle" in backend.framework_name:
+                max_dim = 9
+            for n_axes in range(max_dim):
+                shape = rng.integers(2, 4, size=n_axes)
+                permutation = rng.permutation(n_axes)
+                skipped = 0 if reduction == "rearrange" else rng.integers(n_axes + 1)
+                left = " ".join("x" + str(i) for i in range(n_axes))
+                right = " ".join("x" + str(i) for i in permutation[skipped:])
+                pattern = left + "->" + right
+                x = np.arange(1, 1 + np.prod(shape), dtype=dtype).reshape(shape)
+                if reduction == "prod":
+                    x /= x.mean()  # to avoid overflows
+                result1 = reduce(x, pattern, reduction=reduction)
+                result2 = x.transpose(permutation)
+                if skipped > 0:
+                    result2 = getattr(result2, reduction)(axis=tuple(range(skipped)))
+                assert coincide(result1, result2)
+                check_op_against_numpy(backend, x, pattern, reduction=reduction, axes_lengths={}, is_symbolic=False)
+
+
+def test_reduction_with_callable_imperatives():
+    x_numpy = np.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]).astype("float32")
+    x_numpy /= x_numpy.max()
+
+    def logsumexp_torch(x, tuple_of_axes):
+        return x.logsumexp(tuple_of_axes)
+
+    def logsumexp_tf(x, tuple_of_axes):
+        import tensorflow as tf
+
+        return tf.reduce_logsumexp(x, tuple_of_axes)
+
+    def logsumexp_keras(x, tuple_of_axes):
+        import tensorflow.keras.backend as k
+
+        return k.logsumexp(x, tuple_of_axes)
+
+    def logsumexp_numpy(x, tuple_of_axes):
+        # very naive logsumexp to compare to
+        minused = x.max(tuple_of_axes)
+        y = x - x.max(tuple_of_axes, keepdims=True)
+        y = np.exp(y)
+        y = np.sum(y, axis=tuple_of_axes)
+        return np.log(y) + minused
+
+    from einops._backends import NumpyBackend, TensorflowBackend, TFKerasBackend, TorchBackend
+
+    backend2callback = {
+        TorchBackend.framework_name: logsumexp_torch,
+        TensorflowBackend.framework_name: logsumexp_tf,
+        TFKerasBackend.framework_name: logsumexp_keras,
+        NumpyBackend.framework_name: logsumexp_numpy,
+    }
+
+    for backend in imp_op_backends:
+        if backend.framework_name not in backend2callback:
+            continue
+
+        backend_callback = backend2callback[backend.framework_name]
+
+        x_backend = backend.from_numpy(x_numpy)
+        for pattern1, pattern2 in equivalent_reduction_patterns:
+            print("Test reduction with callable for ", backend.framework_name, pattern1, pattern2)
+            output_numpy = reduce(x_numpy, pattern1, reduction=logsumexp_numpy)
+            output_backend = reduce(x_backend, pattern1, reduction=backend_callback)
+            assert np.allclose(
+                output_numpy,
+                backend.to_numpy(output_backend),
+            )
+
+
+def test_enumerating_directions():
+    for backend in imp_op_backends:
+        print("testing directions for", backend.framework_name)
+        for shape in [[], [1], [1, 1, 1], [2, 3, 5, 7]]:
+            x = np.arange(np.prod(shape)).reshape(shape)
+            axes1 = _enumerate_directions(x)
+            axes2 = _enumerate_directions(backend.from_numpy(x))
+            assert len(axes1) == len(axes2) == len(shape)
+            axes2 = [backend.to_numpy(ax) for ax in axes2]
+            for ax1, ax2 in zip(axes1, axes2):
+                assert ax1.shape == ax2.shape
+                assert np.allclose(ax1, ax2)
+
+
+def test_concatenations_and_stacking():
+    for backend in imp_op_backends:
+        print("testing shapes for ", backend.framework_name)
+        for n_arrays in [1, 2, 5]:
+            shapes = [[], [1], [1, 1], [2, 3, 5, 7], [1] * 6]
+            for shape in shapes:
+                arrays1 = [np.arange(i, i + np.prod(shape)).reshape(shape) for i in range(n_arrays)]
+                arrays2 = [backend.from_numpy(array) for array in arrays1]
+                result0 = np.asarray(arrays1)
+                result1 = rearrange(arrays1, "...->...")
+                result2 = rearrange(arrays2, "...->...")
+                assert np.array_equal(result0, result1)
+                assert np.array_equal(result1, backend.to_numpy(result2))
+
+                result1 = rearrange(arrays1, "b ... -> ... b")
+                result2 = rearrange(arrays2, "b ... -> ... b")
+                assert np.array_equal(result1, backend.to_numpy(result2))
+
+
+def test_gradients_imperatives():
+    # lazy - just checking reductions
+    for reduction in REDUCTIONS:
+        if reduction in ("any", "all"):
+            continue  # non-differentiable ops
+        x = np.arange(1, 1 + 2 * 3 * 4).reshape([2, 3, 4]).astype("float32")
+        results = {}
+        for backend in imp_op_backends:
+            y0 = backend.from_numpy(x)
+            if not hasattr(y0, "grad"):
+                continue
+
+            y1 = reduce(y0, "a b c -> c a", reduction=reduction)
+            y2 = reduce(y1, "c a -> a c", reduction=reduction)
+            y3 = reduce(y2, "a (c1 c2) -> a", reduction=reduction, c1=2)
+            y4 = reduce(y3, "... -> ", reduction=reduction)
+
+            y4.backward()
+            grad = backend.to_numpy(y0.grad)
+            results[backend.framework_name] = grad
+
+        print("comparing gradients for", results.keys())
+        for name1, grad1 in results.items():
+            for name2, grad2 in results.items():
+                assert np.allclose(grad1, grad2), [name1, name2, "provided different gradients"]
+
+
+def test_tiling_imperatives():
+    for backend in imp_op_backends:
+        print("Tiling tests for ", backend.framework_name)
+        input = np.arange(2 * 3 * 5, dtype="int64").reshape([2, 1, 3, 1, 5])
+        test_cases = [
+            (1, 1, 1, 1, 1),
+            (1, 2, 1, 3, 1),
+            (3, 1, 1, 4, 1),
+        ]
+        for repeats in test_cases:
+            expected = np.tile(input, repeats)
+            converted = backend.from_numpy(input)
+            repeated = backend.tile(converted, repeats)
+            result = backend.to_numpy(repeated)
+            assert np.array_equal(result, expected)
+
+
+def test_tiling_symbolic():
+    for backend in sym_op_backends:
+        print("Tiling tests for ", backend.framework_name)
+        input = np.arange(2 * 3 * 5, dtype="int64").reshape([2, 1, 3, 1, 5])
+        test_cases = [
+            (1, 1, 1, 1, 1),
+            (1, 2, 1, 3, 1),
+            (3, 1, 1, 4, 1),
+        ]
+        for repeats in test_cases:
+            expected = np.tile(input, repeats)
+            sym = backend.create_symbol(input.shape)
+            result = backend.eval_symbol(backend.tile(sym, repeats), [[sym, input]])
+            assert np.array_equal(result, expected)
+
+            sym = backend.create_symbol([None] * len(input.shape))
+            result = backend.eval_symbol(backend.tile(sym, repeats), [[sym, input]])
+            assert np.array_equal(result, expected)
+
+
+repeat_test_cases = [
+    # all assume that input has shape [2, 3, 5]
+    ("a b c -> c a b", dict()),
+    ("a b c -> (c copy a b)", dict(copy=2, a=2, b=3, c=5)),
+    ("a b c -> (a copy) b c ", dict(copy=1)),
+    ("a b c -> (c a) (copy1 b copy2)", dict(a=2, copy1=1, copy2=2)),
+    ("a ...  -> a ... copy", dict(copy=4)),
+    ("... c -> ... (copy1 c copy2)", dict(copy1=1, copy2=2)),
+    ("...  -> ... ", dict()),
+    (" ...  -> copy1 ... copy2 ", dict(copy1=2, copy2=3)),
+    ("a b c  -> copy1 a copy2 b c () ", dict(copy1=2, copy2=1)),
+]
+
+
+def check_reversion(x, repeat_pattern, **sizes):
+    """Checks repeat pattern by running reduction"""
+    left, right = repeat_pattern.split("->")
+    reduce_pattern = right + "->" + left
+    repeated = repeat(x, repeat_pattern, **sizes)
+    reduced_min = reduce(repeated, reduce_pattern, reduction="min", **sizes)
+    reduced_max = reduce(repeated, reduce_pattern, reduction="max", **sizes)
+    assert np.array_equal(x, reduced_min)
+    assert np.array_equal(x, reduced_max)
+
+
+def test_repeat_numpy():
+    # check repeat vs reduce. Repeat works ok if reverse reduction with min and max work well
+    x = np.arange(2 * 3 * 5).reshape([2, 3, 5])
+    x1 = repeat(x, "a b c -> copy a b c ", copy=1)
+    assert np.array_equal(x[None], x1)
+    for pattern, axis_dimensions in repeat_test_cases:
+        check_reversion(x, pattern, **axis_dimensions)
+
+
+def test_repeat_imperatives():
+    x = np.arange(2 * 3 * 5).reshape([2, 3, 5])
+    for backend in imp_op_backends:
+        print("Repeat tests for ", backend.framework_name)
+
+        for pattern, axis_dimensions in repeat_test_cases:
+            expected = repeat(x, pattern, **axis_dimensions)
+            converted = backend.from_numpy(x)
+            repeated = repeat(converted, pattern, **axis_dimensions)
+            result = backend.to_numpy(repeated)
+            assert np.array_equal(result, expected)
+
+
+def test_repeat_symbolic():
+    x = np.arange(2 * 3 * 5).reshape([2, 3, 5])
+
+    for backend in sym_op_backends:
+        print("Repeat tests for ", backend.framework_name)
+
+        for pattern, axis_dimensions in repeat_test_cases:
+            expected = repeat(x, pattern, **axis_dimensions)
+
+            sym = backend.create_symbol(x.shape)
+            result = backend.eval_symbol(repeat(sym, pattern, **axis_dimensions), [[sym, x]])
+            assert np.array_equal(result, expected)
+
+
+def test_repeat_array_api():
+    import numpy as xp
+
+    from einops import array_api as AA
+
+    if xp.__version__ < "2.0.0":
+        pytest.skip()
+
+    x = np.arange(2 * 3 * 5).reshape([2, 3, 5])
+
+    for pattern, axis_dimensions in repeat_test_cases:
+        expected = repeat(x, pattern, **axis_dimensions)
+
+        result = AA.repeat(xp.from_dlpack(x), pattern, **axis_dimensions)
+        assert np.array_equal(AA.asnumpy(result + 0), expected)
+
+
+test_cases_repeat_anonymous = [
+    # all assume that input has shape [1, 2, 4, 6]
+    ("a b c d -> c a d b", dict()),
+    ("a b c d -> (c 2 d a b)", dict(a=1, c=4, d=6)),
+    ("1 b c d -> (d copy 1) 3 b c ", dict(copy=3)),
+    ("1 ...  -> 3 ... ", dict()),
+    ("() ... d -> 1 (copy1 d copy2) ... ", dict(copy1=2, copy2=3)),
+    ("1 b c d -> (1 1) (1 b) 2 c 3 d (1 1)", dict()),
+]
+
+
+def test_anonymous_axes():
+    x = np.arange(1 * 2 * 4 * 6).reshape([1, 2, 4, 6])
+    for pattern, axis_dimensions in test_cases_repeat_anonymous:
+        check_reversion(x, pattern, **axis_dimensions)
+
+
+def test_list_inputs():
+    x = np.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
+
+    assert np.array_equal(
+        rearrange(list(x), "... -> (...)"),
+        rearrange(x, "... -> (...)"),
+    )
+    assert np.array_equal(
+        reduce(list(x), "a ... e -> (...)", "min"),
+        reduce(x, "a ... e -> (...)", "min"),
+    )
+    assert np.array_equal(
+        repeat(list(x), "...  -> b (...)", b=3),
+        repeat(x, "...  -> b (...)", b=3),
+    )
+
+
+def test_torch_compile_with_dynamic_shape():
+    if not is_backend_tested("torch"):
+        pytest.skip()
+    import torch
+
+    # somewhat reasonable debug messages
+    torch._dynamo.config.verbose = True
+
+    def func1(x):
+        # test contains ellipsis
+        a, b, c, *other = x.shape
+        x = rearrange(x, "(a a2) b c ... -> b (c a2) (a ...)", a2=2)
+        # test contains passing expression as axis length
+        x = reduce(x, "b ca2 A -> b A", "sum", ca2=c * 2)
+        return x
+
+    # seems can't test static and dynamic in the same test run.
+    func1_compiled_static = torch.compile(func1, dynamic=False, fullgraph=True)
+    func1_compiled_dynamic = torch.compile(func1, dynamic=True, fullgraph=True)
+
+    x = torch.randn(size=[4, 5, 6, 3])
+    assert torch.allclose(func1_compiled_static(x), func1(x), atol=1e-5)
+    assert torch.allclose(func1_compiled_dynamic(x), func1(x), atol=1e-5)
+    # check with input of different dimensionality, and with all shape elements changed
+    x = torch.randn(size=[6, 3, 4, 2, 3])
+    assert torch.allclose(func1_compiled_static(x), func1(x), atol=1e-5)
+    assert torch.allclose(func1_compiled_dynamic(x), func1(x), atol=1e-5)
+
+
+def bit_count(x):
+    return sum((x >> i) & 1 for i in range(20))
+
+
+def test_reduction_imperatives_booleans():
+    """Checks that any/all reduction works in all frameworks"""
+    x_np = np.asarray([(bit_count(x) % 2) == 0 for x in range(2**6)]).reshape([2] * 6)
+    for backend in imp_op_backends:
+        print("Reduction any/all tests for ", backend.framework_name)
+
+        for axis in range(6):
+            expected_result_any = np.any(x_np, axis=axis, keepdims=True)
+            expected_result_all = np.all(x_np, axis=axis, keepdims=True)
+            assert not np.array_equal(expected_result_any, expected_result_all)
+
+            axes = list("abcdef")
+            axes_in = list(axes)
+            axes_out = list(axes)
+            axes_out[axis] = "1"
+            pattern = (" ".join(axes_in)) + " -> " + (" ".join(axes_out))
+
+            res_any = reduce(backend.from_numpy(x_np), pattern, reduction="any")
+            res_all = reduce(backend.from_numpy(x_np), pattern, reduction="all")
+
+            assert np.array_equal(expected_result_any, backend.to_numpy(res_any))
+            assert np.array_equal(expected_result_all, backend.to_numpy(res_all))
+
+        # expected result: any/all
+        expected_result_any = np.any(x_np, axis=(0, 1), keepdims=True)
+        expected_result_all = np.all(x_np, axis=(0, 1), keepdims=True)
+        pattern = "a b ... -> 1 1 ..."
+        res_any = reduce(backend.from_numpy(x_np), pattern, reduction="any")
+        res_all = reduce(backend.from_numpy(x_np), pattern, reduction="all")
+        assert np.array_equal(expected_result_any, backend.to_numpy(res_any))
+        assert np.array_equal(expected_result_all, backend.to_numpy(res_all))

kronos-venv/lib/python3.12/site-packages/einops/tests/test_other.py

@@ -0,0 +1,363 @@
+import subprocess
+import tempfile
+from doctest import testmod
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+import einops
+import einops.layers
+from einops._backends import AbstractBackend
+from einops.einops import _optimize_transformation, parse_shape, rearrange
+from einops.tests import collect_test_backends, is_backend_tested
+
+__author__ = "Alex Rogozhnikov"
+
+rng = np.random.default_rng()
+
+
+def test_doctests_examples():
+    # tests docstrings, additionally
+    testmod(einops.layers, raise_on_error=True, extraglobs=dict(np=np))
+    testmod(einops.einops, raise_on_error=True, extraglobs=dict(np=np))
+
+
+def test_backends_installed():
+    """
+    This test will fail if some of backends are not installed or can't be imported
+    Other tests will just work and only test installed backends.
+    """
+    from . import parse_backends_to_test
+
+    backends_to_test = set(parse_backends_to_test())
+    errors = []
+    # Find backend subclasses recursively
+    backend_subclasses = []
+    backends = AbstractBackend.__subclasses__()
+    while backends:
+        backend = backends.pop()
+        backends += backend.__subclasses__()
+        backend_subclasses.append(backend)
+
+    for backend_type in backend_subclasses:
+        if backend_type.framework_name not in backends_to_test:
+            continue
+        try:
+            # instantiate
+            backend_type()
+            backends_to_test.remove(backend_type.framework_name)
+        except Exception as e:
+            errors.append((backend_type.framework_name, e))
+    assert len(errors) == 0, errors
+    assert len(backends_to_test) == 0, f"did not instantiate {backends_to_test=}, they won't be tested"
+
+
+def test_optimize_transformations_numpy():
+    print("Testing optimizations")
+    shapes = [[2] * n_dimensions for n_dimensions in range(14)]
+    shapes += [[3] * n_dimensions for n_dimensions in range(6)]
+    shapes += [[2, 3, 5, 7]]
+    shapes += [[2, 3, 5, 7, 11, 17]]
+
+    for shape in shapes:
+        for _attempt in range(5):
+            n_dimensions = len(shape)
+            x = rng.integers(0, 2**12, size=shape).reshape([-1])
+            init_shape = shape[:]
+            n_reduced = rng.integers(0, n_dimensions + 1)
+            reduced_axes = tuple(rng.permutation(n_dimensions)[:n_reduced])
+            axes_reordering = rng.permutation(n_dimensions - n_reduced)
+            final_shape = rng.integers(0, 1024, size=333)  # just random
+
+            init_shape2, reduced_axes2, axes_reordering2, final_shape2 = combination2 = _optimize_transformation(
+                init_shape, reduced_axes, axes_reordering, final_shape
+            )
+
+            assert np.array_equal(final_shape, final_shape2)
+            result1 = x.reshape(init_shape).sum(axis=reduced_axes).transpose(axes_reordering).reshape([-1])
+            result2 = x.reshape(init_shape2).sum(axis=reduced_axes2).transpose(axes_reordering2).reshape([-1])
+            assert np.array_equal(result1, result2)
+
+            # testing we can't optimize this formula again
+            combination3 = _optimize_transformation(*combination2)
+            for a, b in zip(combination2, combination3):
+                assert np.array_equal(a, b)
+
+
+_IMPERATIVE_BACKENDS = collect_test_backends(symbolic=False, layers=False)
+
+x_np = np.zeros([10, 20, 30, 40])
+
+
+def test_parse_shape_imperative():
+    for backend in _IMPERATIVE_BACKENDS:
+        print("Shape parsing for ", backend.framework_name)
+        parsed1 = parse_shape(x_np, "a b c d")
+        parsed2 = parse_shape(backend.from_numpy(x_np), "a b c d")
+        assert parsed1 == parsed2 == dict(a=10, b=20, c=30, d=40)
+        assert parsed1 != dict(a=1, b=20, c=30, d=40) != parsed2
+
+
+def test_underscore():
+    for backend in _IMPERATIVE_BACKENDS:
+        parsed1 = parse_shape(x_np, "_ _ _ _")
+        parsed2 = parse_shape(backend.from_numpy(x_np), "_ _ _ _")
+        assert parsed1 == parsed2 == dict()
+
+
+def test_underscore_one():
+    for backend in _IMPERATIVE_BACKENDS:
+        parsed1 = parse_shape(x_np, "_ _ _ hello")
+        parsed2 = parse_shape(backend.from_numpy(x_np), "_ _ _ hello")
+        assert parsed1 == parsed2 == dict(hello=40)
+
+
+def test_underscore_several():
+    for backend in _IMPERATIVE_BACKENDS:
+        parsed1 = parse_shape(x_np, "_ _ a1 a1a111a")
+        parsed2 = parse_shape(backend.from_numpy(x_np), "_ _ a1 a1a111a")
+        assert parsed1 == parsed2 == dict(a1=30, a1a111a=40)
+
+
+def test_repeating():
+    with pytest.raises(einops.EinopsError):
+        parse_shape(x_np, "a a b b")
+
+    for backend in _IMPERATIVE_BACKENDS:
+        with pytest.raises(einops.EinopsError):
+            parse_shape(backend.from_numpy(x_np), "a a b b")
+
+
+def test_ellipsis():
+    for backend in _IMPERATIVE_BACKENDS:
+        for shape, pattern, expected in [
+            ([10, 20], "...", dict()),
+            ([10], "... a", dict(a=10)),
+            ([10, 20], "... a", dict(a=20)),
+            ([10, 20, 30], "... a", dict(a=30)),
+            ([10, 20, 30, 40], "... a", dict(a=40)),
+            ([10], "a ...", dict(a=10)),
+            ([10, 20], "a ...", dict(a=10)),
+            ([10, 20, 30], "a ...", dict(a=10)),
+            ([10, 20, 30, 40], "a ...", dict(a=10)),
+            ([10, 20, 30, 40], " a ... b", dict(a=10, b=40)),
+            ([10, 40], " a ... b", dict(a=10, b=40)),
+        ]:
+            x = np.ones(shape)
+            parsed1 = parse_shape(x, pattern)
+            parsed2 = parse_shape(backend.from_numpy(x), pattern)
+            assert parsed1 == parsed2 == expected
+
+
+def test_parse_with_anonymous_axes():
+    for backend in _IMPERATIVE_BACKENDS:
+        for shape, pattern, expected in [
+            ([1, 2, 3, 4], "1 2 3 a", dict(a=4)),
+            ([10, 1, 2], "a 1 2", dict(a=10)),
+            ([10, 1, 2], "a () 2", dict(a=10)),
+        ]:
+            x = np.ones(shape)
+            parsed1 = parse_shape(x, pattern)
+            parsed2 = parse_shape(backend.from_numpy(x), pattern)
+            assert parsed1 == parsed2 == expected
+
+
+def test_failures():
+    for backend in _IMPERATIVE_BACKENDS:
+        # every test should fail
+        for shape, pattern in [
+            ([1, 2, 3, 4], "a b c"),
+            ([1, 2, 3, 4], "2 a b c"),
+            ([1, 2, 3, 4], "a b c ()"),
+            ([1, 2, 3, 4], "a b c d e"),
+            ([1, 2, 3, 4], "a b c d e ..."),
+            ([1, 2, 3, 4], "a b c ()"),
+        ]:
+            with pytest.raises(RuntimeError):
+                x = np.ones(shape)
+                parse_shape(backend.from_numpy(x), pattern)
+
+
+_SYMBOLIC_BACKENDS = [
+    *collect_test_backends(symbolic=True, layers=False),
+    *collect_test_backends(symbolic=True, layers=True),
+]
+
+# tensorflow.keras needs special way to compile,
+# shape vars can be used only inside layers but not as outputs
+_SYMBOLIC_BACKENDS = [backend for backend in _SYMBOLIC_BACKENDS if backend.framework_name != "tensorflow.keras"]
+
+
+@pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS)
+def test_parse_shape_symbolic(backend):
+    for input_shape in [
+        [10, 20, 30, 40],
+        [10, 20, None, None],
+        [None, None, None, None],
+    ]:
+        print(f"special shape parsing {backend.framework_name=} {input_shape=}")
+        input_symbol = backend.create_symbol(input_shape)
+
+        shape_placeholder = parse_shape(input_symbol, "a b c d")
+        out_shape = {}
+        for name, symbol in shape_placeholder.items():
+            out_shape[name] = (
+                symbol
+                if isinstance(symbol, int)
+                else backend.eval_symbol(symbol, [(input_symbol, np.zeros([10, 20, 30, 40]))])
+            )  # out shape element is either int, or symbol that we are able to eval
+        print(out_shape)
+        result_placeholder = rearrange(
+            input_symbol, "a b (c1 c2) (d1 d2) -> (a b d1) c1 (c2 d2)", **parse_shape(input_symbol, "a b c1 _"), d2=2
+        )
+        result = backend.eval_symbol(result_placeholder, [(input_symbol, np.zeros([10, 20, 30, 40]))])
+        print(result.shape)
+        assert result.shape == (10 * 20 * 20, 30, 1 * 2)
+        assert np.allclose(result, 0)
+
+
+@pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS)
+def test_parse_shape_symbolic_ellipsis(backend):
+    for static_shape, shape, pattern, expected in [
+        ([10, 20], [None, None], "...", dict()),
+        ([10], [None], "... a", dict(a=10)),
+        ([10, 20], [None, None], "... a", dict(a=20)),
+        ([10, 20, 30], [None, None, None], "... a", dict(a=30)),
+        ([10, 20, 30, 40], [None, None, None, None], "... a", dict(a=40)),
+        ([10], [None], "a ...", dict(a=10)),
+        ([10, 20], [None, None], "a ...", dict(a=10)),
+        ([10, 20, 30], [None, None, None], "a ...", dict(a=10)),
+        ([10, 20, 30, 40], [None, None, None, None], "a ...", dict(a=10)),
+        ([10, 20, 30, 40], [None, None, None, None], " a ... b", dict(a=10, b=40)),
+        ([10, 40], [None, None], " a ... b ", dict(a=10, b=40)),
+    ]:
+        input_symbol = backend.create_symbol(shape)
+        shape_placeholder = parse_shape(input_symbol, pattern)
+        out_shape = {}
+        for name, symbol in shape_placeholder.items():
+            if isinstance(symbol, int):
+                out_shape[name] = symbol
+            else:
+                out_shape[name] = backend.eval_symbol(symbol, [(input_symbol, np.zeros(static_shape))])
+        assert out_shape == expected
+
+
+def test_is_float_type():
+    backends = collect_test_backends(symbolic=False, layers=False)
+    backends += collect_test_backends(symbolic=False, layers=True)
+    for backend in backends:
+        for dtype in ["int32", "int64", "float32", "float64"]:
+            is_float = "float" in dtype
+            input = np.zeros([3, 4, 5], dtype=dtype)
+            input = backend.from_numpy(input)
+            assert backend.is_float_type(input) == is_float, (dtype, backend, input.dtype)
+
+
+def test_torch_compile_for_functions():
+    """
+    Test ensures that allow_ops_in_compiled_graph allows compiling in a single graph
+    Additionally we ensure that after compilation cache works properly
+     (by changing shapes and patterns)
+    We additionally check that pack/unpack still can be handled
+     despite variable number of inputs/outputs
+    """
+    if not is_backend_tested("torch"):
+        pytest.skip()
+    import torch
+    from torch import nn
+
+    from einops import einsum, pack, reduce, repeat, unpack
+    from einops._torch_specific import allow_ops_in_compiled_graph
+
+    allow_ops_in_compiled_graph()
+
+    class TorchModuleWithOperations(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+
+        def forward(self, x_abc, suffix=""):
+            a, b, c = x_abc.shape
+
+            def suf(pattern):
+                parts = pattern.split()
+                return " ".join([p if p[-1] not in "acd" else p + suffix for p in parts])
+
+            # patterns look a bit strange because names a, c, d will be modified on every run
+            # by suf function
+            x_abcd = repeat(x_abc, suf("a b c -> a b c 4"))
+            x_abc = reduce(x_abcd, suf("a b c d -> a b c"), "min")
+            x_abdc, ps = pack([x_abc] * (2 + len(suffix)), suf("a b * c"))
+            x_array = unpack(rearrange(x_abdc, suf("a b d c -> (a b ) 1 c d")), ps, "ab one1 c *")
+            x1 = x_array[0] + len(x_array)
+            x1 = rearrange(x1, suf("(a b ) 1 c -> a b c"), b=b)
+            addition = einsum(x_abc, x_abcd, suf("a b c , a b c d -> d"))[0]
+            return x1 + addition
+
+    original = TorchModuleWithOperations()
+    compiled = torch.compile(original, fullgraph=True)
+    for size in [10, 20, 40]:
+        x = torch.rand([size, size + 1, size + 2])
+        for suffix in ["", "suf1", "other_suffix"]:
+            result1 = compiled(x, suffix)
+            result2 = original(x.double(), suffix).float()
+
+            torch.testing.assert_close(result1, result2, atol=1e-5, rtol=1e-5)
+
+
+def test_torch_compile_for_layers():
+    """
+    Einops layers are in general very friendly towards tracing/compiling,
+    but we still want to make sure we can compile them.
+    """
+    if not is_backend_tested("torch"):
+        pytest.skip()
+
+    import torch
+    from torch import nn
+
+    from einops.layers.torch import EinMix, Rearrange, Reduce
+
+    original = nn.Sequential(
+        Rearrange("b (t c) -> b t c", c=16),
+        EinMix("b t c -> qkv b t cout", weight_shape="qkv c cout", bias_shape="qkv cout", qkv=3, c=16, cout=8),
+        Reduce("qkv b t cout -> b t qkv", "min", cout=8),
+    )
+
+    compiled = torch.compile(original, fullgraph=True)
+
+    for size in [16, 32, 64]:
+        x = torch.rand([size, size])
+        result1 = original(x)
+        result2 = compiled(x)
+        assert torch.allclose(result1, result2)
+
+
+src = """
+import einops
+import numpy as np
+from concurrent.futures import ThreadPoolExecutor
+import torch
+
+def f():
+    return einops.rearrange(np.ndarray((20, 150, 150)), "... i j -> ... j i")
+with ThreadPoolExecutor(max_workers=2) as ex:
+    fs = []
+    for i in range(20):
+        fs.append(ex.submit(f))
+    for fut in fs:
+        fut.result()
+"""
+
+
+def test_einops_threading():
+    # requires both. Reproduces problem from https://github.com/arogozhnikov/einops/issues/391
+    if not is_backend_tested("torch"):
+        pytest.skip()
+    if not is_backend_tested("numpy"):
+        pytest.skip()
+
+    with tempfile.TemporaryDirectory() as d:
+        testfile = Path(d).joinpath("test.py")
+        testfile.write_text(src)
+        subprocess.run(["python", testfile.absolute().as_posix()], check=True)

kronos-venv/lib/python3.12/site-packages/einops/tests/test_packing.py

@@ -0,0 +1,312 @@
+import dataclasses
+import typing
+
+import numpy as np
+import pytest
+
+from einops import EinopsError, asnumpy, pack, unpack
+from einops.tests import collect_test_backends
+
+rng = np.random.default_rng()
+
+
+def pack_unpack(xs, pattern):
+    x, ps = pack(xs, pattern)
+    unpacked = unpack(xs, ps, pattern)
+    assert len(unpacked) == len(xs)
+    for a, b in zip(unpacked, xs):
+        assert np.allclose(asnumpy(a), asnumpy(b))
+
+
+def unpack_and_pack(x, ps, pattern: str):
+    unpacked = unpack(x, ps, pattern)
+    packed, ps2 = pack(unpacked, pattern=pattern)
+
+    assert np.allclose(asnumpy(packed), asnumpy(x))
+    return unpacked
+
+
+def unpack_and_pack_against_numpy(x, ps, pattern: str):
+    capturer_backend = CaptureException()
+    capturer_numpy = CaptureException()
+
+    with capturer_backend:
+        unpacked = unpack(x, ps, pattern)
+        packed, ps2 = pack(unpacked, pattern=pattern)
+
+    with capturer_numpy:
+        x_np = asnumpy(x)
+        unpacked_np = unpack(x_np, ps, pattern)
+        packed_np, ps3 = pack(unpacked_np, pattern=pattern)
+
+    assert type(capturer_numpy.exception) == type(capturer_backend.exception)  # noqa E721
+    if capturer_numpy.exception is not None:
+        # both failed
+        return
+    else:
+        # neither failed, check results are identical
+        assert np.allclose(asnumpy(packed), asnumpy(x))
+        assert np.allclose(asnumpy(packed_np), asnumpy(x))
+        assert len(unpacked) == len(unpacked_np)
+        for a, b in zip(unpacked, unpacked_np):
+            assert np.allclose(asnumpy(a), b)
+
+
+class CaptureException:
+    def __enter__(self):
+        self.exception = None
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.exception = exc_val
+        return True
+
+
+def test_numpy_trivial(H=13, W=17):
+    def rand(*shape):
+        return rng.random(shape)
+
+    def check(a, b):
+        assert a.dtype == b.dtype
+        assert a.shape == b.shape
+        assert np.all(a == b)
+
+    r, g, b = rand(3, H, W)
+    embeddings = rand(H, W, 32)
+
+    check(
+        np.stack([r, g, b], axis=2),
+        pack([r, g, b], "h w *")[0],
+    )
+    check(
+        np.stack([r, g, b], axis=1),
+        pack([r, g, b], "h * w")[0],
+    )
+    check(
+        np.stack([r, g, b], axis=0),
+        pack([r, g, b], "* h w")[0],
+    )
+
+    check(
+        np.concatenate([r, g, b], axis=1),
+        pack([r, g, b], "h *")[0],
+    )
+    check(
+        np.concatenate([r, g, b], axis=0),
+        pack([r, g, b], "* w")[0],
+    )
+
+    i = np.index_exp[:, :, None]
+    check(
+        np.concatenate([r[i], g[i], b[i], embeddings], axis=2),
+        pack([r, g, b, embeddings], "h w *")[0],
+    )
+
+    with pytest.raises(EinopsError):
+        pack([r, g, b, embeddings], "h w nonexisting_axis *")
+
+    pack([r, g, b], "some_name_for_H some_name_for_w1 *")
+
+    with pytest.raises(EinopsError):
+        pack([r, g, b, embeddings], "h _w *")  # no leading underscore
+    with pytest.raises(EinopsError):
+        pack([r, g, b, embeddings], "h_ w *")  # no trailing underscore
+    with pytest.raises(EinopsError):
+        pack([r, g, b, embeddings], "1h_ w *")
+    with pytest.raises(EinopsError):
+        pack([r, g, b, embeddings], "1 w *")
+    with pytest.raises(EinopsError):
+        pack([r, g, b, embeddings], "h h *")
+    # capital and non-capital are different
+    pack([r, g, b, embeddings], "h H *")
+
+
+@dataclasses.dataclass
+class UnpackTestCase:
+    shape: typing.Tuple[int, ...]
+    pattern: str
+
+    def dim(self):
+        return self.pattern.split().index("*")
+
+    def selfcheck(self):
+        assert self.shape[self.dim()] == 5
+
+
+cases = [
+    # NB: in all cases unpacked axis is of length 5.
+    # that's actively used in tests below
+    UnpackTestCase((5,), "*"),
+    UnpackTestCase((5, 7), "* seven"),
+    UnpackTestCase((7, 5), "seven *"),
+    UnpackTestCase((5, 3, 4), "* three four"),
+    UnpackTestCase((4, 5, 3), "four * three"),
+    UnpackTestCase((3, 4, 5), "three four *"),
+]
+
+
+def test_pack_unpack_with_numpy():
+    case: UnpackTestCase
+
+    for case in cases:
+        shape = case.shape
+        pattern = case.pattern
+
+        x = rng.random(shape)
+        # all correct, no minus 1
+        unpack_and_pack(x, [[2], [1], [2]], pattern)
+        # no -1, asking for wrong shapes
+        with pytest.raises(EinopsError):
+            unpack_and_pack(x, [[2], [1], [2]], pattern + " non_existent_axis")
+        with pytest.raises(EinopsError):
+            unpack_and_pack(x, [[2], [1], [1]], pattern)
+        with pytest.raises(EinopsError):
+            unpack_and_pack(x, [[4], [1], [1]], pattern)
+        # all correct, with -1
+        unpack_and_pack(x, [[2], [1], [-1]], pattern)
+        unpack_and_pack(x, [[2], [-1], [2]], pattern)
+        unpack_and_pack(x, [[-1], [1], [2]], pattern)
+        _, _, last = unpack_and_pack(x, [[2], [3], [-1]], pattern)
+        assert last.shape[case.dim()] == 0
+        # asking for more elements than available
+        with pytest.raises(EinopsError):
+            unpack(x, [[2], [4], [-1]], pattern)
+        # this one does not raise, because indexing x[2:1] just returns zero elements
+        # with pytest.raises(EinopsError):
+        #     unpack(x, [[2], [-1], [4]], pattern)
+        with pytest.raises(EinopsError):
+            unpack(x, [[-1], [1], [5]], pattern)
+
+        # all correct, -1 nested
+        rs = unpack_and_pack(x, [[1, 2], [1, 1], [-1, 1]], pattern)
+        assert all(len(r.shape) == len(x.shape) + 1 for r in rs)
+        rs = unpack_and_pack(x, [[1, 2], [1, -1], [1, 1]], pattern)
+        assert all(len(r.shape) == len(x.shape) + 1 for r in rs)
+        rs = unpack_and_pack(x, [[2, -1], [1, 2], [1, 1]], pattern)
+        assert all(len(r.shape) == len(x.shape) + 1 for r in rs)
+
+        # asking for more elements, -1 nested
+        with pytest.raises(EinopsError):
+            unpack(x, [[-1, 2], [1], [5]], pattern)
+        with pytest.raises(EinopsError):
+            unpack(x, [[2, 2], [2], [5, -1]], pattern)
+
+        # asking for non-divisible number of elements
+        with pytest.raises(EinopsError):
+            unpack(x, [[2, 1], [1], [3, -1]], pattern)
+        with pytest.raises(EinopsError):
+            unpack(x, [[2, 1], [3, -1], [1]], pattern)
+        with pytest.raises(EinopsError):
+            unpack(x, [[3, -1], [2, 1], [1]], pattern)
+
+        # -1 takes zero
+        unpack_and_pack(x, [[0], [5], [-1]], pattern)
+        unpack_and_pack(x, [[0], [-1], [5]], pattern)
+        unpack_and_pack(x, [[-1], [5], [0]], pattern)
+
+        # -1 takes zero, -1
+        unpack_and_pack(x, [[2, -1], [1, 5]], pattern)
+
+
+def test_pack_unpack_against_numpy():
+    for backend in collect_test_backends(symbolic=False, layers=False):
+        print(f"test packing against numpy for {backend.framework_name}")
+        check_zero_len = True
+
+        for case in cases:
+            unpack_and_pack = unpack_and_pack_against_numpy
+            shape = case.shape
+            pattern = case.pattern
+
+            x = rng.random(shape)
+            x = backend.from_numpy(x)
+            # all correct, no minus 1
+            unpack_and_pack(x, [[2], [1], [2]], pattern)
+            # no -1, asking for wrong shapes
+            with pytest.raises(EinopsError):
+                unpack(x, [[2], [1], [1]], pattern)
+
+            with pytest.raises(EinopsError):
+                unpack(x, [[4], [1], [1]], pattern)
+            # all correct, with -1
+            unpack_and_pack(x, [[2], [1], [-1]], pattern)
+            unpack_and_pack(x, [[2], [-1], [2]], pattern)
+            unpack_and_pack(x, [[-1], [1], [2]], pattern)
+
+            # asking for more elements than available
+            with pytest.raises(EinopsError):
+                unpack(x, [[2], [4], [-1]], pattern)
+            # this one does not raise, because indexing x[2:1] just returns zero elements
+            # with pytest.raises(EinopsError):
+            #     unpack(x, [[2], [-1], [4]], pattern)
+            with pytest.raises(EinopsError):
+                unpack(x, [[-1], [1], [5]], pattern)
+
+            # all correct, -1 nested
+            unpack_and_pack(x, [[1, 2], [1, 1], [-1, 1]], pattern)
+            unpack_and_pack(x, [[1, 2], [1, -1], [1, 1]], pattern)
+            unpack_and_pack(x, [[2, -1], [1, 2], [1, 1]], pattern)
+
+            # asking for more elements, -1 nested
+            with pytest.raises(EinopsError):
+                unpack(x, [[-1, 2], [1], [5]], pattern)
+            with pytest.raises(EinopsError):
+                unpack(x, [[2, 2], [2], [5, -1]], pattern)
+
+            # asking for non-divisible number of elements
+            with pytest.raises(EinopsError):
+                unpack(x, [[2, 1], [1], [3, -1]], pattern)
+            with pytest.raises(EinopsError):
+                unpack(x, [[2, 1], [3, -1], [1]], pattern)
+            with pytest.raises(EinopsError):
+                unpack(x, [[3, -1], [2, 1], [1]], pattern)
+
+            if check_zero_len:
+                # -1 takes zero
+                unpack_and_pack(x, [[2], [3], [-1]], pattern)
+                unpack_and_pack(x, [[0], [5], [-1]], pattern)
+                unpack_and_pack(x, [[0], [-1], [5]], pattern)
+                unpack_and_pack(x, [[-1], [5], [0]], pattern)
+
+                # -1 takes zero, -1
+                unpack_and_pack(x, [[2, -1], [1, 5]], pattern)
+
+
+def test_pack_unpack_array_api():
+    import numpy as xp
+
+    from einops import array_api as AA
+
+    if xp.__version__ < "2.0.0":
+        pytest.skip()
+
+    for case in cases:
+        shape = case.shape
+        pattern = case.pattern
+        x_np = rng.random(shape)
+        x_xp = xp.from_dlpack(x_np)
+
+        for ps in [
+            [[2], [1], [2]],
+            [[1], [1], [-1]],
+            [[1], [1], [-1, 3]],
+            [[2, 1], [1, 1, 1], [-1]],
+        ]:
+            x_np_split = unpack(x_np, ps, pattern)
+            x_xp_split = AA.unpack(x_xp, ps, pattern)
+            for a, b in zip(x_np_split, x_xp_split):
+                assert np.allclose(a, AA.asnumpy(b + 0))
+
+            x_agg_np, ps1 = pack(x_np_split, pattern)
+            x_agg_xp, ps2 = AA.pack(x_xp_split, pattern)
+            assert ps1 == ps2
+            assert np.allclose(x_agg_np, AA.asnumpy(x_agg_xp))
+
+        for ps in [
+            [[2, 3]],
+            [[1], [5]],
+            [[1], [5], [-1]],
+            [[1], [2, 3]],
+            [[1], [5], [-1, 2]],
+        ]:
+            with pytest.raises(EinopsError):
+                unpack(x_np, ps, pattern)

kronos-venv/lib/python3.12/site-packages/einops/tests/test_parsing.py

@@ -0,0 +1,126 @@
+import pytest
+
+from einops import EinopsError
+from einops.parsing import AnonymousAxis, ParsedExpression, _ellipsis
+
+__author__ = "Alex Rogozhnikov"
+
+
+class AnonymousAxisPlaceholder:
+    def __init__(self, value: int):
+        self.value = value
+        assert isinstance(self.value, int)
+
+    def __eq__(self, other):
+        return isinstance(other, AnonymousAxis) and self.value == other.value
+
+
+def test_anonymous_axes():
+    a, b = AnonymousAxis("2"), AnonymousAxis("2")
+    assert a != b
+    c, d = AnonymousAxisPlaceholder(2), AnonymousAxisPlaceholder(3)
+    assert a == c and b == c
+    assert a != d and b != d
+    assert [a, 2, b] == [c, 2, c]
+
+
+def test_elementary_axis_name():
+    for name in [
+        "a",
+        "b",
+        "h",
+        "dx",
+        "h1",
+        "zz",
+        "i9123",
+        "somelongname",
+        "Alex",
+        "camelCase",
+        "u_n_d_e_r_score",
+        "unreasonablyLongAxisName",
+    ]:
+        assert ParsedExpression.check_axis_name(name)
+
+    for name in ["", "2b", "12", "_startWithUnderscore", "endWithUnderscore_", "_", "...", _ellipsis]:
+        assert not ParsedExpression.check_axis_name(name)
+
+
+def test_invalid_expressions():
+    # double ellipsis should raise an error
+    ParsedExpression("... a b c d")
+    with pytest.raises(EinopsError):
+        ParsedExpression("... a b c d ...")
+    with pytest.raises(EinopsError):
+        ParsedExpression("... a b c (d ...)")
+    with pytest.raises(EinopsError):
+        ParsedExpression("(... a) b c (d ...)")
+
+    # double/missing/enclosed parenthesis
+    ParsedExpression("(a) b c (d ...)")
+    with pytest.raises(EinopsError):
+        ParsedExpression("(a)) b c (d ...)")
+    with pytest.raises(EinopsError):
+        ParsedExpression("(a b c (d ...)")
+    with pytest.raises(EinopsError):
+        ParsedExpression("(a) (()) b c (d ...)")
+    with pytest.raises(EinopsError):
+        ParsedExpression("(a) ((b c) (d ...))")
+
+    # invalid identifiers
+    ParsedExpression("camelCase under_scored cApiTaLs ß ...")
+    with pytest.raises(EinopsError):
+        ParsedExpression("1a")
+    with pytest.raises(EinopsError):
+        ParsedExpression("_pre")
+    with pytest.raises(EinopsError):
+        ParsedExpression("...pre")
+    with pytest.raises(EinopsError):
+        ParsedExpression("pre...")
+
+
+def test_parse_expression():
+    parsed = ParsedExpression("a1  b1   c1    d1")
+    assert parsed.identifiers == {"a1", "b1", "c1", "d1"}
+    assert parsed.composition == [["a1"], ["b1"], ["c1"], ["d1"]]
+    assert not parsed.has_non_unitary_anonymous_axes
+    assert not parsed.has_ellipsis
+
+    parsed = ParsedExpression("() () () ()")
+    assert parsed.identifiers == set()
+    assert parsed.composition == [[], [], [], []]
+    assert not parsed.has_non_unitary_anonymous_axes
+    assert not parsed.has_ellipsis
+
+    parsed = ParsedExpression("1 1 1 ()")
+    assert parsed.identifiers == set()
+    assert parsed.composition == [[], [], [], []]
+    assert not parsed.has_non_unitary_anonymous_axes
+    assert not parsed.has_ellipsis
+
+    aap = AnonymousAxisPlaceholder
+
+    parsed = ParsedExpression("5 (3 4)")
+    assert len(parsed.identifiers) == 3 and {i.value for i in parsed.identifiers} == {3, 4, 5}
+    assert parsed.composition == [[aap(5)], [aap(3), aap(4)]]
+    assert parsed.has_non_unitary_anonymous_axes
+    assert not parsed.has_ellipsis
+
+    parsed = ParsedExpression("5 1 (1 4) 1")
+    assert len(parsed.identifiers) == 2 and {i.value for i in parsed.identifiers} == {4, 5}
+    assert parsed.composition == [[aap(5)], [], [aap(4)], []]
+
+    parsed = ParsedExpression("name1 ... a1 12 (name2 14)")
+    assert len(parsed.identifiers) == 6
+    assert parsed.identifiers.difference({"name1", _ellipsis, "a1", "name2"}).__len__() == 2
+    assert parsed.composition == [["name1"], _ellipsis, ["a1"], [aap(12)], ["name2", aap(14)]]
+    assert parsed.has_non_unitary_anonymous_axes
+    assert parsed.has_ellipsis
+    assert not parsed.has_ellipsis_parenthesized
+
+    parsed = ParsedExpression("(name1 ... a1 12) name2 14")
+    assert len(parsed.identifiers) == 6
+    assert parsed.identifiers.difference({"name1", _ellipsis, "a1", "name2"}).__len__() == 2
+    assert parsed.composition == [["name1", _ellipsis, "a1", aap(12)], ["name2"], [aap(14)]]
+    assert parsed.has_non_unitary_anonymous_axes
+    assert parsed.has_ellipsis
+    assert parsed.has_ellipsis_parenthesized