flax EinMix 语法与 torch 不同

在 torch 中，EinMix 实现使用 `**axes_lengths` 来初始化轴长，而 flax 实现需要提供一个 `sizes` 字典。这导致两个框架所需的代码不同，但如果 einops 在所有其他模块中保持统一，那将是很棒的。