Einops库新增匿名轴解析功能，提升张量形状检查便捷性

Einops是一个强大的张量操作库，最近在其parse_shape函数中新增了对匿名轴的支持，这一改进显著提升了代码的可读性和编写效率。

功能背景

在深度学习开发中，经常需要检查张量的形状是否符合预期。Einops库提供了parse_shape函数，可以将张量形状与指定的模式字符串进行匹配。例如：

parse_shape(torch.tensor([[0]]), "b n")  # 返回 {'b': 1, 'n': 1}

然而，之前的版本不支持直接在模式字符串中使用具体数字作为匿名轴，这导致开发者需要为每个固定尺寸的轴创建命名变量，增加了代码复杂度。

新增功能解析

最新版本的Einops现在支持以下用法：

parse_shape(torch.tensor([[0]]), "b 1")  # 现在可以正常工作

这一改进使得开发者可以直接在模式字符串中使用具体数字表示固定尺寸的轴，而无需为每个固定尺寸创建命名变量。这与Einops中已有的下划线(_)语法类似，后者用于表示"忽略此轴"。

实际应用场景

这个功能特别适用于以下场景：

1. 形状断言检查：在模型开发中验证中间张量的形状
2. API文档：清晰地表达函数输入输出的形状要求
3. 测试代码：简化张量形状的验证逻辑

例如，现在可以这样编写形状检查代码：

def assert_shapes(*tensor_patterns, **dim_sizes):
    shapes = [
        dict(**dim_sizes),
        *(parse_shape(tensor, pattern) for tensor, pattern in tensor_patterns),
    ]
    
    for dim_name in set(dim for shape in shapes for dim in shape):
        sizes = set(shape.get(dim_name, None) for shape in shapes)
        sizes.discard(None)
        assert len(sizes) <= 1, f"尺寸'{dim_name}'不一致: {sizes}"

# 更简洁的形状检查
assert_shapes(
    (tensor1, 'batch 1 h'),
    (tensor2, 'batch 2 h'),
    h=64
)

技术实现原理

在底层实现上，Einops的parse_shape函数现在会：

1. 解析模式字符串时，将数字视为匿名轴
2. 验证对应轴的实际尺寸是否匹配指定数字
3. 在返回的字典中不包含这些匿名轴的信息

这与处理下划线(_)的机制类似，但增加了尺寸验证的功能。

对开发者的意义

这一改进使得：

1. 代码更简洁：减少不必要的变量命名
2. 意图更清晰：直接看到预期的固定尺寸
3. 维护更方便：减少因变量名变更导致的修改

特别是在处理具有多个固定尺寸轴的张量时，这一功能可以显著减少样板代码。

总结

Einops库通过支持匿名轴解析，进一步提升了其在张量操作领域的实用性和便利性。这一改进虽然看似微小，但在实际开发中能带来明显的效率提升和代码质量改进，体现了Einops团队对开发者体验的持续关注。