einops

einops — это библиотека Python с открытым исходным кодом, предоставляющая гибкие и мощные операции с тензорами через нотацию, вдохновленную Эйнштейном, предназначенную для повышения читаемости и надежности кода манипулирования тензорами. Название расшифровывается как Einstein-Inspired Notation for operations, частично вдохновлено numpy.einsum. Основной API минималистичен, но мощен, и построен вокруг трех основных функций: rearrange (перестановка элементов по шаблону), reduce (объединение перестановки с редукцией, такой как mean или max) и repeat (копирование по новой оси). Позднее были добавлены pack и unpack для обратимого объединения нескольких тензоров разной размерности, а также einsum с поддержкой многосимвольных имен осей. Einops также включает слой EinMix — универсальный линейный слой, хорошо подходящий для MLP Mixers и аналогичных архитектур, а также обертки слоев для конкретных фреймворков (Rearrange, Reduce и др.) для PyTorch, TensorFlow, Flax и Paddle. Нотация подчеркивает семантическую информацию: вместо написания x.view(x.shape[0], -1) пользователи пишут rearrange(x, 'b c h w -> b (c h w)'), что документирует как входные, так и выходные формы. Именованные оси, такие как 'time c h w -> time (c h w)', указывают, что данные представляют видеопоследовательность, а не независимый батч. Пользователи могут также закреплять конкретные размеры (например, c=256, h=19, w=19), чтобы операция проверяла формы во время выполнения, а не полагалась на комментарии. Результаты строго детерминированы, поэтому неоднозначные операции, такие как depth-to-space, можно записывать однозначно. Тот же синтаксис единообразно работает для пулинга 1D, 2D и 3D, а поведение согласовано между фреймворками, даже если нативные API различаются (например, семантика flatten отличается между NumPy/PyTorch и TensorFlow/Keras). Einops поддерживает NumPy, PyTorch, TensorFlow, JAX, CuPy, Flax, Paddle, OneFlow, tinygrad и PyTensor, а также работает с любым фреймворком, реализующим стандарт массива Python API, включая NumPy >= 2.0, MLX, pydata/sparse, cubed и nndonnx. Поддерживаются torch.compile и torch.jit.script для слоев PyTorch. Библиотека требует Python 3.10 или новее и устанавливается через pip install einops. Проект был представлен как устный доклад на ICLR 2022 и используется более чем в 10 000 проектов на GitHub.