einops

此库受到 Python 的 einops 的极大启发。

目前 tch 是唯一可用的后端。

与 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 和 prod 等操作来归约轴，如果需要在多个连续轴上执行相同的操作，我们可以这样做 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")