形状和布局

XLA Shape proto (xla_data.proto) 描述了 N 维数组(简称“数组”)的秩、大小和数据类型。

术语、表示法和惯例

  • 数组的排位等于维度的数量。数组的真实秩是指大小大于 1 的维度的数量。

  • 对于 N 维数组,维度从 0N-1 进行编号。为方便起见,维度编号为任意标签。这些尺寸数字的顺序并不意味着形状布局中的特定次要/主要顺序。布局由 Layout proto 决定。

  • 按照惯例,维度按维度编号递增顺序列出。例如,对于大小为 [A x B x C] 的三维数组,维度 0 的大小为 A,维度 1 的大小为 B,维度 2 的大小为 C

    XLA 中的某些实用程序还支持类似 Python 的否定索引:维度 -1 是最后一个维度(等同于 N 维度数组的 N-1)。例如,对于上述三维数组,维度 -1 的大小为 C,维度 -2 的大小为 B,依此类推。

  • 二维、三维和四维数组通常具有与维度关联的特定字母。例如,对于二维数组:

    • 维度 0:y
    • 维度 1:x

    对于 3D 数组:

    • 维度 0:z
    • 维度 1:y
    • 维度 2:x

    对于四维数组:

    • 维度 0:p
    • 维度 1:z
    • 维度 2:y
    • 维度 3:x
  • XLA API 中接受维度的函数会按维度编号的递增顺序执行此操作。这与将维度作为 initializer_list 传递时使用的顺序一致;例如:

    ShapeUtil::MakeShape(F32, {A, B, C, D})

    将创建一个尺寸大小数组由序列 [A, B, C, D] 组成的形状。

布局

Layout proto 描述了数组在内存中的表示方式。Layout proto 包含以下字段:

message Layout {
  repeated int64 minor_to_major = 1;
  repeated int64 padded_dimensions = 2;
  optional PaddingValue padding_value = 3;
}

从次要维度到主要维度排序

唯一的必填字段是 minor_to_major。此字段描述了形状内尺寸的从小到大的顺序。minor_to_major 中的值是数组的维度排序(N 维度数组为 0N-1),第一个值从最小维度到最后一个值,即最大维度。最小维度是指在线性内存布局的数组元素中单步调试时变化最快的维度。

例如,请考虑以下大小为 [2 x 3] 的二维数组:

a b c
d e f

此处,维度 0 为尺寸 2,维度 1 为尺寸 3。如果布局中的 minor_to_major 字段为 [0, 1],则维度 0 是最次要维度,维度 1 是最主要维度。这对应于线性内存中的以下布局:

a d b e c f

这个从次要维度到主要维度顺序(从 0N-1)类似于“主要列”(排名为 2)。假设维度的单调顺序,我们可以在代码中指代此布局的另一种方式是“dim 0 是次要的”。

另一方面,如果布局中的 minor_to_major 字段为 [1, 0],则线性内存中的布局为:

a b c d e f

对于 N 维度数组,从次要维度到主要维度顺序(从 N-10)类似于行主要维度(排名为 2)。假设维度的单调顺序,我们可以在代码中引用此布局的另一种方式是“dim 0 is main”。

默认的从次要到主要排序

新创建的形状的默认布局为“维度顺序为主要顺序”(类似于 2 阶的“主要”行)。

内边距

内边距在可选的 padded_dimensionspadding_value 字段中定义。padded_dimensions 字段描述填充每个维度的尺寸(宽度)。如果存在,padded_dimensions 中的元素数量必须等于形状的阶。

例如,给定上面定义的 [2 x 3] 数组,如果 padded_dimensions[3, 5],维度 0 将填充为宽度 3,维度 1 填充为宽度 5。线性内存中的布局(假设内边距值为 0,主布局)为:

a d 0 b e 0 c f 0 0 0 0 0 0 0

这相当于以下数组的布局,这些数组具有相同的次要到主要维度顺序:

a b c 0 0
d e f 0 0
0 0 0 0 0

编入数组索引

index_util.h 中的 IndexUtil 类提供了在给定形状和布局的多维索引和线性索引之间进行转换的实用程序。多维索引包含每个维度的 int64 索引。线性索引是单个 int64 值,它会编入包含数组的缓冲区中的索引。如需了解可用于简化形状和布局创建和操作的实用程序,请参阅同一目录中的 shape_util.hlayout_util.h