candle-einops

此库是从 einops 分支出来的，旨在将 einops 的支持引入 Candle。感谢 @VasanthakumarV 提供如此出色的基于宏的库。原始库使用 TCH 作为后端库，并基于 einops Python 库实现。

原始库的大部分内容都保留了下来，只是改变了设备/数据类型绑定。我计划将来将 einsum 功能移植到这个库中。

与 Python 版本的区别

• 编译时生成所有代码，避免了缓存的需要
• 一个通用的 API 用于重新排列、归约和重复操作
• 可以在表达式中直接指定形状和归约操作

入门指南

转置

交换/转置维度，-> 的左侧是原始状态，右侧 -> 描述的是最终状态

// (28, 28, 3) becomes (3, 28, 28)
let output = einops!("h w c -> c h w", &input);

组合

通过在 -> 右侧括号内放置维度来组合维度

// (10, 28, 28, 3) becomes (280, 28, 3)
let output = einops!("b h w c -> (b h) w c", &input);

转置 + 组合

在一个表达式中执行转置操作，然后组合两个维度为一个维度

// (10, 28, 28, 3) becomes (28, 280, 3)
let output = einops!("b h w c -> h (b w) c", &input);

分解

通过在左侧括号内指定详细信息将一个维度拆分为两个，例如指定新维度的形状为 b1:2，b1 是一个形状为 2 的新维度

// (10, 28, 28, 3) becomes (2, 5, 28, 28, 3)
let output = einops!("(b1:2 b2) h w c -> b1 b2 h w c", &input);

也可以使用大括号从变量或字段（结构体和枚举）中指定新轴

let b1 = 2;
let output = einops!("({b1} b2) h w c -> {b1} b2 h w c", &input);

分解 + 转置 + 组合

我们可以在单个表达式中执行到目前为止讨论的所有操作

// (10, 28, 28, 3) becomes (56, 140 3)
let output = einops!("b h (w w2:2) c -> (h w2) (b w) c", &input);

归约

我们可以使用诸如 sum、min、max 和 mean 等操作来减少轴。如果需要在多个连续轴上执行相同的操作，我们可以使用 sum(a b c)

// (10, 28, 28, 3) becomes (28, 28, 3)
let output = einops!("mean(b) h w c -> h w c", &input);

分解 + 减少 + 转置 + 组合

一个表达式可以组合所有讨论的功能

// (10, 28, 28, 3) becomes (14, 140, 3)
let output = einops!("b (h max(h2:2)) (w max(w2:2)) c -> h (b w) c", &input);

重复

我们可以在 -> 的右侧指定重复轴，它可以有名称，或者它也可以简单地是一个数字

// (28, 28, 3) becomes (28, 5, 28, 3)
let output = einops!("h w c -> h repeat:5 w c", &input);

重复轴的形状可以来自变量或字段（结构体、枚举）

let repeat = 5;
let output = einops!("h w c -> h {repeat} w c", &input);

挤压

挤压形状为 1 的轴

// (1, 28, 28, 3) becomes (28, 28, 3)
let output = einops!("1 h w c -> h w c")